fix: import unsloth before transformers/trl to avoid lazy-import bug

training/training_unsloth.py

@@ -42,6 +42,17 @@ def _build_args() -> argparse.Namespace:
     parser.add_argument("--model_name", default="unsloth/Qwen2.5-3B-Instruct")
     parser.add_argument("--scenario", default=None)
     parser.add_argument("--difficulty", choices=["easy", "medium", "hard"], default="easy")
+    parser.add_argument(
+        "--curriculum",
+        action="store_true",
+        help=(
+            "Enable adaptive curriculum: start at --difficulty and promote "
+            "to medium/hard once held-out success rate clears the threshold "
+            "(see training/curriculum.py)."
+        ),
+    )
+    parser.add_argument("--curriculum_promote", type=float, default=0.55)
+    parser.add_argument("--curriculum_demote", type=float, default=0.10)
     parser.add_argument("--total_episodes", type=int, default=400)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--max_steps", type=int, default=18)
@@ -68,27 +79,55 @@ def _build_args() -> argparse.Namespace:
 def main() -> None:  # pragma: no cover - heavy GPU path
     args = _build_args()
 
-    from datasets import Dataset
+    # IMPORTANT: Unsloth MUST be imported before transformers / trl. It
+    # patches transformers' lazy ``_import_structure`` to register a few
+    # symbols (notably ``PreTrainedModel`` under torch-aware paths). If trl
+    # loads transformers first, the lazy loader will fail with a confusing
+    # ``ImportError: cannot import name 'PreTrainedModel' from 'transformers'``
+    # at GRPOTrainer import time — which is exactly what we hit on the
+    # trainer Space before this reorder.
+    # See: https://github.com/unslothai/unsloth and the matching
+    # transformers issue #42548 for the lazy-import root cause.
+    from unsloth import FastLanguageModel
     from transformers import TrainerCallback
     from trl import GRPOConfig, GRPOTrainer
-    from unsloth import FastLanguageModel
 
     from server.environment import CERNCollisionEnvironment
+    from training.curriculum import CurriculumConfig, CurriculumManager
     from training.evidence import (
         CheckpointEvalWriter,
         EvidencePaths,
+        RewardComponentLogWriter,
         TrainingLogWriter,
         render_checkpoint_progression,
+        render_reward_components,
         render_training_curve,
     )
-    from training.llm_agent import LLMAgentConfig, build_chat
+    from training.llm_agent import LLMAgentConfig
     from training.rollouts import collect_episode
-    from training.training_script import EpisodeContext, _format_validity_bonus, _stepwise_reward
+    from training.training_script import (
+        EpisodeContext,
+        RewardComponentAccumulator,
+    )
 
     paths = EvidencePaths(root=Path(args.evidence_dir))
     paths.ensure()
     log_writer = TrainingLogWriter(paths.training_log_csv)
     ckpt_writer = CheckpointEvalWriter(paths.checkpoint_evals_csv)
+    component_writer = RewardComponentLogWriter(paths.reward_components_csv)
+    component_accumulator = RewardComponentAccumulator()
+
+    curriculum: Optional[CurriculumManager] = None
+    if args.curriculum:
+        curriculum = CurriculumManager(
+            CurriculumConfig(
+                start_difficulty=args.difficulty,
+                promote_threshold=args.curriculum_promote,
+                demote_threshold=args.curriculum_demote,
+            )
+        )
+        logger.info("Curriculum enabled: start=%s promote≥%.2f demote≤%.2f",
+                    args.difficulty, args.curriculum_promote, args.curriculum_demote)
 
     logger.info("Loading Unsloth model: %s", args.model_name)
     model, tokenizer = FastLanguageModel.from_pretrained(
@@ -110,28 +149,23 @@ def main() -> None:  # pragma: no cover - heavy GPU path
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
+    from training.training_script import build_dataset, make_reward_fn
+
     env = CERNCollisionEnvironment(max_steps=args.max_steps)
-    prompts: List[str] = []
-    for i in range(args.total_episodes):
-        obs = env.reset(seed=args.seed + i, scenario=args.scenario, difficulty=args.difficulty)
-        chat = build_chat(obs)
-        prompts.append(
-            tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=False)
-        )
-    dataset = Dataset.from_dict({"prompt": prompts})
+    dataset = build_dataset(
+        tokenizer=tokenizer,
+        n_prompts=args.total_episodes,
+        seed=args.seed,
+        scenario=args.scenario,
+        difficulty=args.difficulty,
+        curriculum=args.curriculum,
+    )
 
     ctx = EpisodeContext(
         env=env, seed=args.seed,
         scenario=args.scenario, difficulty=args.difficulty,
     )
-
-    def reward_fn(prompts: List[str], completions: List[str], **kwargs: Any) -> List[float]:
-        rewards: List[float] = []
-        for completion in completions:
-            r = _stepwise_reward(completion_text=completion, ctx=ctx)
-            r += _format_validity_bonus(completion)
-            rewards.append(float(r))
-        return rewards
+    reward_fn = make_reward_fn(ctx, accumulator=component_accumulator)
 
     cfg = GRPOConfig(
         output_dir=args.output_dir,
@@ -174,6 +208,18 @@ def main() -> None:  # pragma: no cover - heavy GPU path
                 log_writer.append(row)
                 render_training_curve(paths.training_log_csv, paths.training_curve_png)
 
+            # Per-component reward summary (FAQ Q17, Q43, Q52: don't watch
+            # only the mean reward — track terminal vs shaping, success
+            # rates, and parse rate so verifier hacks become visible).
+            drained = component_accumulator.drain()
+            if drained:
+                summary = RewardComponentAccumulator.summarise(drained)
+                summary["step"] = state.global_step
+                component_writer.append(summary)
+                render_reward_components(
+                    paths.reward_components_csv, paths.reward_components_png,
+                )
+
         def on_step_end(self, _args, state, control, **kw):
             step = state.global_step
             if step <= 0 or step == self._last_eval_step:
@@ -190,20 +236,29 @@ def main() -> None:  # pragma: no cover - heavy GPU path
         def _run_checkpoint_eval(self, step: int, state) -> None:
             FastLanguageModel.for_inference(model)
             try:
+                # When curriculum is enabled, evaluate at whatever tier the
+                # adaptive manager currently considers appropriate. Otherwise
+                # use the static --difficulty.
+                eval_difficulty = (
+                    curriculum.next_difficulty()
+                    if curriculum is not None
+                    else args.difficulty
+                )
                 episodes = []
                 for s in held_out_seeds:
-                    ep = self._rollout_one(seed=s)
+                    ep = self._rollout_one(seed=s, difficulty=eval_difficulty)
                     if ep is not None:
                         episodes.append(ep)
                 if not episodes:
                     return
                 rewards = [e.cumulative_reward for e in episodes]
+                success_rate = sum(1 for e in episodes if e.discovered) / len(episodes)
                 ckpt_writer.append(
                     step=step,
                     fraction_done=round(step / max(state.max_steps or step, 1), 4),
                     episodes=len(episodes),
                     mean_reward=round(sum(rewards) / len(rewards), 4),
-                    success_rate=round(sum(1 for e in episodes if e.discovered) / len(episodes), 4),
+                    success_rate=round(success_rate, 4),
                     mass_acc=round(sum(1 for e in episodes if e.correct_mass) / len(episodes), 4),
                     channel_acc=round(sum(1 for e in episodes if e.correct_channel) / len(episodes), 4),
                 )
@@ -211,15 +266,27 @@ def main() -> None:  # pragma: no cover - heavy GPU path
                     paths.checkpoint_evals_csv,
                     paths.checkpoint_progression_png,
                 )
+                if curriculum is not None:
+                    snap = curriculum.record(
+                        success=success_rate >= 0.5,
+                        reward=sum(rewards) / len(rewards),
+                    )
+                    curriculum.save(paths.root / "curriculum_state.json")
+                    if snap.get("event"):
+                        logger.info(
+                            "[curriculum] %s @ step=%d → tier=%s (rolling=%.2f)",
+                            snap["event"], step, snap["current"], snap["rolling_success"],
+                        )
                 logger.info(
-                    "[checkpoint-eval step=%d] reward=%.3f success=%.2f",
-                    step, rewards and (sum(rewards) / len(rewards)) or 0.0,
-                    sum(1 for e in episodes if e.discovered) / len(episodes),
+                    "[checkpoint-eval step=%d difficulty=%s] reward=%.3f success=%.2f",
+                    step, eval_difficulty,
+                    rewards and (sum(rewards) / len(rewards)) or 0.0,
+                    success_rate,
                 )
             finally:
                 FastLanguageModel.for_training(model)
 
-        def _rollout_one(self, seed: int):
+        def _rollout_one(self, seed: int, difficulty: Optional[str] = None):
             def prompt_fn(chat):
                 return tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=False)
 
@@ -236,7 +303,8 @@ def main() -> None:  # pragma: no cover - heavy GPU path
 
             return collect_episode(
                 env=env, seed=seed,
-                scenario=args.scenario, difficulty=args.difficulty,
+                scenario=args.scenario,
+                difficulty=difficulty or args.difficulty,
                 prompt_fn=prompt_fn, generate_fn=generate_fn,
                 config=LLMAgentConfig(),
             )
@@ -251,6 +319,18 @@ def main() -> None:  # pragma: no cover - heavy GPU path
     )
     logger.info("Starting Unsloth + LoRA GRPO training")
     trainer.train()
+
+    # Drain whatever rollouts the final on_log didn't catch so the last
+    # row of reward_components.csv is correct.
+    final_drain = component_accumulator.drain()
+    if final_drain:
+        summary = RewardComponentAccumulator.summarise(final_drain)
+        summary["step"] = trainer.state.global_step
+        component_writer.append(summary)
+        render_reward_components(
+            paths.reward_components_csv, paths.reward_components_png,
+        )
+
     trainer.save_model(args.output_dir)
     tokenizer.save_pretrained(args.output_dir)
     logger.info("Saved adapters to %s", args.output_dir)