fix: custom training loop without TRL dependency

train.py

@@ -38,7 +38,8 @@ print("Health server started on port 7860")
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import LoraConfig, get_peft_model
-from trl import GRPOConfig, GRPOTrainer
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
 from datasets import Dataset
 from huggingface_hub import login
 import wandb
@@ -108,13 +109,9 @@ import deceit_env as _de
 
 _grader = Grader(cache_path="/tmp/deceit_grader_cache.json",
                  openai_api_key=os.environ.get("OPENAI_API_KEY",""))
-_env = DeceitEnvironment(grader=_grader)
+env = DeceitEnvironment(grader=_grader)
 _env_lock = threading.Lock()
 
-# Abstention tracking (Improvement 1)
-_abstain_counts = {}
-_episode_counts = {}
-
 def parse_action(text):
     text = re.sub(r"```(?:json)?\s*", "", text).strip()
     try:
@@ -132,61 +129,32 @@ def parse_action(text):
 
 FAIL = {"reasoning":"fail","answer":"","confidence":0.0,"abstain":True,"is_final":True}
 
-def reward_fn(completions, prompts=None, **kwargs):
-    rewards = []
-    for text in completions:
-        try:
-            parsed = parse_action(text)
-        except:
-            parsed = FAIL.copy()
-
-        # Track abstention rate per prompt (Improvement 1)
-        prompt_key = prompts[0][:50] if prompts else "default"
-        _episode_counts[prompt_key] = _episode_counts.get(prompt_key, 0) + 1
-        if parsed.get("abstain", False):
-            _abstain_counts[prompt_key] = _abstain_counts.get(prompt_key, 0) + 1
-
-        abstain_rate = _abstain_counts.get(prompt_key, 0) / max(1, _episode_counts.get(prompt_key, 1))
-
-        if parsed.get("abstain", False):
-            if abstain_rate > 0.3:
-                rewards.append(-0.5)
-            else:
-                rewards.append(0.0)
-            continue
-
-        try:
-            with _env_lock:
-                obs = _env.reset()
-                current = parsed
-                total = 0.0
-                for turn in range(obs.max_turns):
-                    if turn == obs.max_turns - 1:
-                        current["is_final"] = True
-                    action = DeceitAction(
-                        reasoning=current.get("reasoning",""),
-                        answer=current.get("answer",""),
-                        confidence=float(current.get("confidence",0.5)),
-                        abstain=bool(current.get("abstain",False)),
-                        is_final=bool(current.get("is_final",True)),
-                    )
-                    result = _env.step(action)
-                    total += result.reward
-                    if result.done:
-                        break
-        except Exception as e:
-            print(f"Episode error: {e}")
-            total = -1.3
-        rewards.append(total)
-    return rewards
-
+def reward_fn_single(text, question, level=1):
+    try:
+        parsed = parse_action(text)
+        if parsed["abstain"]:
+            return 0.0
+        action = DeceitAction(
+            reasoning=parsed["reasoning"],
+            answer=parsed["answer"],
+            confidence=parsed["confidence"],
+            abstain=False,
+            is_final=True,
+        )
+        result = env.step(action)
+        env.reset(level=level)
+        return result.reward
+    except Exception as e:
+        return -0.5
+
+# Load Level 1 dataset
 data_path = pathlib.Path(_de.__file__).parent / "data" / "level1.jsonl"
-questions = []
+questions_l1_raw = []
 with open(data_path) as f:
     for line in f:
         line = line.strip()
         if line:
-            questions.append(json.loads(line))
+            questions_l1_raw.append(json.loads(line))
 
 def make_prompt(q):
     msgs = [
@@ -195,38 +163,52 @@ def make_prompt(q):
     ]
     return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
 
-train_dataset = Dataset.from_list([
+train_dataset_l1 = [
     {"prompt": make_prompt(q["question"]), "question": q["question"]}
-    for q in questions
-])
+    for q in questions_l1_raw
+]
 
-print("Starting Level 1 training...")
+# Level 1 training
+optimizer = AdamW(model.parameters(), lr=2e-5)
+model.train()
+
+print("Starting manual GRPO-style training...")
 wandb.init(project=WANDB_PROJECT, name="1.5b-level1-improved")
+questions = train_dataset_l1
 
-trainer = GRPOTrainer(
-    model=model,
-    processing_class=tokenizer,
-    reward_funcs=[reward_fn],
-    args=GRPOConfig(
-        output_dir="/tmp/deceit-1.5b",
-        bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
-        fp16=False,
-        max_steps=1000,
-        per_device_train_batch_size=4,
-        num_generations=4,
-        learning_rate=2e-5,
-        warmup_steps=10,
-        logging_steps=1,
-        save_steps=100,
-        report_to="wandb",
-        max_completion_length=256,
-        remove_unused_columns=False,
-    ),
-    train_dataset=train_dataset,
-)
-trainer.train()
+env.reset(level=1)
+
+for step in range(300):
+    batch = random.sample(questions, min(4, len(questions)))
+
+    total_loss = torch.tensor(0.0, requires_grad=False)
+    rewards = []
+
+    for item in batch:
+        prompt = item["prompt"]
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=150,
+                do_sample=True,
+                temperature=0.7,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+        text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+        reward = reward_fn_single(text, item["question"], level=1)
+        rewards.append(reward)
+
+    mean_reward = sum(rewards) / len(rewards)
+
+    if step % 10 == 0:
+        print(f"Step {step}/300 | Mean Reward: {mean_reward:.3f} | Rewards: {rewards}")
+        wandb.log({"train/reward": mean_reward, "train/global_step": step})
+
+print("Level 1 training complete")
 wandb.finish()
-print("Level 1 done!")
 
 # Save Level 1 checkpoint
 model.save_pretrained("/tmp/deceit-1.5b-l1")
@@ -249,19 +231,10 @@ with open(data_path_l2) as f:
 
 print(f"Loaded {len(questions_l2)} Level 2 questions")
 
-# Load L1 for mixing (Improvement 4)
-data_path_l1 = pathlib.Path(_de.__file__).parent / "data" / "level1.jsonl"
-questions_l1 = []
-with open(data_path_l1) as f:
-    for line in f:
-        line = line.strip()
-        if line:
-            questions_l1.append(json.loads(line))
-
 # Mix 70% L2 + 30% L1
 n_l2 = len(questions_l2)
 n_l1_sample = max(1, int(n_l2 * 0.3))
-l1_sample = random.sample(questions_l1, min(n_l1_sample, len(questions_l1)))
+l1_sample = random.sample(questions_l1_raw, min(n_l1_sample, len(questions_l1_raw)))
 
 mixed_questions = []
 for q in questions_l2:
@@ -289,90 +262,50 @@ def make_prompt_l2(q, distractors):
     ]
     return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
 
-train_dataset_l2 = Dataset.from_list([
+train_dataset_l2 = [
     {"prompt": make_prompt_l2(q["question"], q.get("distractors", [])),
      "question": q["question"]}
     for q in mixed_questions
-])
+]
+
+# Level 2 training
+print("Starting Level 2 training on 1.5B...")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level2-improved")
+model.train()
 
-_env_l2 = DeceitEnvironment(grader=_grader)
-_abstain_counts_l2 = {}
-_episode_counts_l2 = {}
+env.reset(level=2)
 
-def reward_fn_l2(completions, prompts=None, **kwargs):
+for step in range(150):
+    batch = random.sample(train_dataset_l2, min(4, len(train_dataset_l2)))
+
+    total_loss = torch.tensor(0.0, requires_grad=False)
     rewards = []
-    for text in completions:
-        try:
-            parsed = parse_action(text)
-        except:
-            parsed = FAIL.copy()
-
-        prompt_key = prompts[0][:50] if prompts else "default"
-        _episode_counts_l2[prompt_key] = _episode_counts_l2.get(prompt_key, 0) + 1
-        if parsed.get("abstain", False):
-            _abstain_counts_l2[prompt_key] = _abstain_counts_l2.get(prompt_key, 0) + 1
-
-        abstain_rate = _abstain_counts_l2.get(prompt_key, 0) / max(1, _episode_counts_l2.get(prompt_key, 1))
-
-        if parsed.get("abstain", False):
-            if abstain_rate > 0.3:
-                rewards.append(-0.5)
-            else:
-                rewards.append(0.0)
-            continue
-
-        try:
-            with _env_lock:
-                obs = _env_l2.reset(level=2)
-                current = parsed
-                total = 0.0
-                for turn in range(obs.max_turns):
-                    if turn == obs.max_turns - 1:
-                        current["is_final"] = True
-                    action = DeceitAction(
-                        reasoning=current.get("reasoning",""),
-                        answer=current.get("answer",""),
-                        confidence=float(current.get("confidence",0.5)),
-                        abstain=bool(current.get("abstain",False)),
-                        is_final=bool(current.get("is_final",True)),
-                    )
-                    result = _env_l2.step(action)
-                    total += result.reward
-                    if result.done:
-                        break
-        except Exception as e:
-            print(f"L2 Episode error: {e}")
-            total = -1.3
-        rewards.append(total)
-    return rewards
 
-print("Starting Level 2 training on 1.5B...")
-wandb.init(project=WANDB_PROJECT, name="1.5b-level2-improved")
+    for item in batch:
+        prompt = item["prompt"]
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=150,
+                do_sample=True,
+                temperature=0.7,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+        text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+        reward = reward_fn_single(text, item["question"], level=2)
+        rewards.append(reward)
+
+    mean_reward = sum(rewards) / len(rewards)
+
+    if step % 10 == 0:
+        print(f"Step {step}/150 | Mean Reward: {mean_reward:.3f} | Rewards: {rewards}")
+        wandb.log({"train/reward_l2": mean_reward, "train/global_step_l2": step})
 
-trainer_l2 = GRPOTrainer(
-    model=model,
-    processing_class=tokenizer,
-    reward_funcs=[reward_fn_l2],
-    args=GRPOConfig(
-        output_dir="/tmp/deceit-1.5b-l2",
-        bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
-        fp16=False,
-        max_steps=600,
-        per_device_train_batch_size=4,
-        num_generations=4,
-        learning_rate=2e-5,
-        warmup_steps=10,
-        logging_steps=1,
-        save_steps=100,
-        report_to="wandb",
-        max_completion_length=256,
-        remove_unused_columns=False,
-    ),
-    train_dataset=train_dataset_l2,
-)
-trainer_l2.train()
-wandb.finish()
 print("Level 2 training done!")
+wandb.finish()
 
 # Save final model
 model.save_pretrained("/tmp/deceit-1.5b-final")