add logs and fix train error

scripts/test_space_api.py

@@ -32,6 +32,10 @@ def main() -> None:
     assert "completed_steps" in train_status.json()
     assert "remaining_steps" in train_status.json()
     assert "phase" in train_status.json()
+    assert "run_manifest_path" in train_status.json()
+    assert "events_path" in train_status.json()
+    assert "latest_checkpoint_path" in train_status.json()
+    assert "logs_deleted_from_space" in train_status.json()
 
     reset = client.post("/reset", json={"difficulty": "easy", "problem_id": "sum_even_numbers"})
     assert reset.status_code == 200

scripts/test_trace_logging.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from training.trace_logging import TraceArtifactLogger
+
+
+def main() -> None:
+    with TemporaryDirectory() as tmpdir:
+        output_dir = Path(tmpdir)
+        logger = TraceArtifactLogger(
+            run_id="run-123",
+            output_dir=output_dir,
+            training_config={"max_steps": 6, "model_name": "demo-model"},
+            model_identifiers={"model_name": "demo-model", "generator_mode": "reward_aware"},
+            system_prompt="You are the Solver Agent.",
+            checkpoint_interval_steps=2,
+        )
+
+        manifest_path = output_dir / "logs" / "run_manifest.json"
+        assert manifest_path.exists()
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+        assert manifest["run_id"] == "run-123"
+        assert manifest["training_config"]["max_steps"] == 6
+
+        logger.log_event(
+            {
+                "phase": "train",
+                "step": 0,
+                "train_episode_index": 1,
+                "problem_id": "sum_even_numbers_1",
+                "problem_family": "sum_even_numbers",
+                "difficulty": "easy",
+                "teacher_prompt": "Problem: sum the even numbers",
+                "solver_completion": "print(sum(x for x in nums if x % 2 == 0))",
+                "extracted_code": "print(sum(x for x in nums if x % 2 == 0))",
+                "reward": 0.94,
+                "pass_rate": 1.0,
+                "visible_pass_rate": 1.0,
+                "execution_status": "completed",
+                "efficiency_score": 0.94,
+                "optimization_hints": ["Avoid materializing temporary containers."],
+                "feedback": "All hidden tests passed, but the solution can still be optimized further.",
+            }
+        )
+        logger.record_progress(
+            {
+                "phase": "train",
+                "completed_steps": 2,
+                "total_steps": 6,
+                "remaining_steps": 4,
+                "progress_ratio": 0.3333,
+                "current_epoch": 2.0,
+                "current_difficulty": "easy",
+                "curriculum_level": 1,
+                "train_episode_index": 1,
+                "last_problem_id": "sum_even_numbers_1",
+                "last_problem_family": "sum_even_numbers",
+                "last_execution_status": "completed",
+            }
+        )
+        artifact_paths = logger.artifact_paths()
+
+        events_path = Path(artifact_paths["events_path"])
+        latest_checkpoint_path = Path(artifact_paths["latest_checkpoint_path"])
+        assert events_path.exists()
+        assert latest_checkpoint_path.exists()
+
+        event_line = events_path.read_text(encoding="utf-8").strip().splitlines()[0]
+        event = json.loads(event_line)
+        assert event["problem_id"] == "sum_even_numbers_1"
+        assert event["teacher_prompt"] == "Problem: sum the even numbers"
+        assert "training_config" not in event
+
+        checkpoint = json.loads(latest_checkpoint_path.read_text(encoding="utf-8"))
+        assert checkpoint["step"] == 2
+        assert checkpoint["rolling_metrics"]["avg_reward"] == 0.94
+        assert "training_config" not in checkpoint
+
+        reward_curve = output_dir / "reward_curve.csv"
+        reward_curve.write_text("step,episode_reward\n0,0.94\n", encoding="utf-8")
+        summary_paths = logger.finalize(
+            reward_curve_csv=reward_curve,
+            final_metrics={"completed_steps": 6},
+        )
+        summary_path = Path(summary_paths)
+        assert summary_path.exists()
+        summary = json.loads(summary_path.read_text(encoding="utf-8"))
+        assert summary["final_metrics"]["completed_steps"] == 6
+
+    print("Trace logging smoke tests passed")
+
+
+if __name__ == "__main__":
+    main()

server/runtime.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import json
 import os
+import shutil
 import threading
 import traceback
 from dataclasses import asdict, dataclass, field
@@ -69,6 +70,13 @@ class TrainingJobState:
     reward_curve_csv: str | None = None
     model_repo_id: str | None = None
     uploaded_revision: str | None = None
+    logs_dir: str | None = None
+    run_manifest_path: str | None = None
+    events_path: str | None = None
+    latest_checkpoint_path: str | None = None
+    run_summary_path: str | None = None
+    checkpoint_paths: list[str] = field(default_factory=list)
+    logs_deleted_from_space: bool = False
     phase: str = "idle"
     completed_steps: int = 0
     total_steps: int = 0
@@ -430,6 +438,13 @@ class SpaceTrainingManager:
                 reward_curve_csv=payload.get("reward_curve_csv"),
                 model_repo_id=payload.get("model_repo_id"),
                 uploaded_revision=payload.get("uploaded_revision"),
+                logs_dir=payload.get("logs_dir"),
+                run_manifest_path=payload.get("run_manifest_path"),
+                events_path=payload.get("events_path"),
+                latest_checkpoint_path=payload.get("latest_checkpoint_path"),
+                run_summary_path=payload.get("run_summary_path"),
+                checkpoint_paths=payload.get("checkpoint_paths", []),
+                logs_deleted_from_space=bool(payload.get("logs_deleted_from_space", False)),
                 phase=payload.get("phase", "idle"),
                 completed_steps=int(payload.get("completed_steps", 0) or 0),
                 total_steps=int(payload.get("total_steps", 0) or 0),
@@ -508,6 +523,7 @@ class SpaceTrainingManager:
             else:
                 output_dir = self.output_root / requested_output_dir / run_id
             config.output_dir = str(output_dir)
+            logs_dir = output_dir / "logs"
 
             self._job = TrainingJobState(
                 status="running",
@@ -519,6 +535,12 @@ class SpaceTrainingManager:
                 reward_curve_csv=None,
                 model_repo_id=os.getenv("HF_MODEL_REPO_ID"),
                 uploaded_revision=None,
+                logs_dir=str(logs_dir),
+                run_manifest_path=str(logs_dir / "run_manifest.json"),
+                events_path=str(logs_dir / "events.jsonl"),
+                latest_checkpoint_path=str(logs_dir / "latest_checkpoint.json"),
+                run_summary_path=str(logs_dir / "run_summary.json"),
+                checkpoint_paths=[],
                 phase="queued",
                 completed_steps=0,
                 total_steps=int(config.max_steps),
@@ -562,11 +584,22 @@ class SpaceTrainingManager:
         )
         return getattr(commit_info, "oid", None) or getattr(commit_info, "commit_hash", None) or "unknown"
 
+    def _cleanup_local_logs(self, log_dir: str | None) -> bool:
+        if not log_dir:
+            return False
+        folder_path = Path(log_dir)
+        if not folder_path.exists():
+            return False
+        shutil.rmtree(folder_path, ignore_errors=True)
+        return not folder_path.exists()
+
     def _run_training_job(self, run_id: str, config: TrainingConfig) -> None:
+        summary: dict[str, Any] | None = None
         try:
-            summary = run_training(config, progress_callback=self._update_progress)
+            summary = run_training(config, run_id=run_id, progress_callback=self._update_progress)
             artifact_path = summary["output_dir"]
             uploaded_revision = self._upload_artifacts(artifact_path, run_id)
+            logs_deleted = self._cleanup_local_logs(summary.get("logs_dir"))
             self.model_registry.load_latest_from_hub()
 
             with self._lock:
@@ -576,6 +609,13 @@ class SpaceTrainingManager:
                 self._job.reward_curve_csv = summary.get("reward_curve_csv")
                 self._job.model_repo_id = os.getenv("HF_MODEL_REPO_ID")
                 self._job.uploaded_revision = uploaded_revision
+                self._job.logs_dir = None if logs_deleted else summary.get("logs_dir")
+                self._job.run_manifest_path = None if logs_deleted else summary.get("run_manifest_path")
+                self._job.events_path = None if logs_deleted else summary.get("events_path")
+                self._job.latest_checkpoint_path = None if logs_deleted else summary.get("latest_checkpoint_path")
+                self._job.run_summary_path = None if logs_deleted else summary.get("run_summary_path")
+                self._job.checkpoint_paths = [] if logs_deleted else summary.get("checkpoint_paths", [])
+                self._job.logs_deleted_from_space = logs_deleted
                 self._job.phase = "completed"
                 self._job.completed_steps = int(summary.get("completed_steps", config.max_steps))
                 self._job.total_steps = int(config.max_steps)
@@ -589,9 +629,18 @@ class SpaceTrainingManager:
                 self._job.traceback = None
                 self._persist_status()
         except Exception as exc:
+            logs_deleted = self._cleanup_local_logs(summary.get("logs_dir") if summary else self._job.logs_dir)
             with self._lock:
                 self._job.status = "failed"
                 self._job.finished_at = _utc_now()
+                if logs_deleted:
+                    self._job.logs_dir = None
+                    self._job.run_manifest_path = None
+                    self._job.events_path = None
+                    self._job.latest_checkpoint_path = None
+                    self._job.run_summary_path = None
+                    self._job.checkpoint_paths = []
+                self._job.logs_deleted_from_space = logs_deleted
                 self._job.error = str(exc)
                 self._job.traceback = traceback.format_exc()
                 self._persist_status()

test.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from scripts.test_env import main as run_env_smoke
 from scripts.test_space_api import main as run_space_api_smoke
+from scripts.test_trace_logging import main as run_trace_logging_smoke
 from scripts.test_verifier import test_cases
 from verifier.verifier import verify
 
@@ -9,6 +10,7 @@ from verifier.verifier import verify
 def main() -> None:
     run_env_smoke()
     run_space_api_smoke()
+    run_trace_logging_smoke()
 
     reward, info = verify(
         "n=int(input())\nnums=list(map(int,input().split()))\nprint(sum(x for x in nums if x % 2 == 0))",

training/trace_logging.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import json
+import shutil
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+def _utc_now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _json_safe(value: Any) -> Any:
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, dict):
+        return {str(key): _json_safe(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_json_safe(item) for item in value]
+    return value
+
+
+@dataclass
+class TraceArtifactLogger:
+    run_id: str
+    output_dir: Path
+    training_config: dict[str, Any]
+    model_identifiers: dict[str, Any]
+    system_prompt: str
+    checkpoint_interval_steps: int = 10
+    schema_version: str = "1.0"
+    logs_dir: Path = field(init=False)
+    manifest_path: Path = field(init=False)
+    events_path: Path = field(init=False)
+    latest_checkpoint_path: Path = field(init=False)
+    run_summary_path: Path = field(init=False)
+    checkpoint_paths: list[Path] = field(default_factory=list, init=False)
+    _last_checkpoint_step: int = field(default=0, init=False, repr=False)
+    _latest_event: dict[str, Any] = field(default_factory=dict, init=False, repr=False)
+    _recent_rewards: deque[float] = field(default_factory=lambda: deque(maxlen=25), init=False, repr=False)
+    _recent_pass_rates: deque[float] = field(default_factory=lambda: deque(maxlen=25), init=False, repr=False)
+    _recent_efficiency_scores: deque[float] = field(default_factory=lambda: deque(maxlen=25), init=False, repr=False)
+    _latest_progress: dict[str, Any] = field(default_factory=dict, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self.logs_dir = self.output_dir / "logs"
+        self.logs_dir.mkdir(parents=True, exist_ok=True)
+        self.manifest_path = self.logs_dir / "run_manifest.json"
+        self.events_path = self.logs_dir / "events.jsonl"
+        self.latest_checkpoint_path = self.logs_dir / "latest_checkpoint.json"
+        self.run_summary_path = self.logs_dir / "run_summary.json"
+        manifest = {
+            "run_id": self.run_id,
+            "schema_version": self.schema_version,
+            "started_at": _utc_now_iso(),
+            "training_config": self.training_config,
+            "model_identifiers": self.model_identifiers,
+            "system_prompt": self.system_prompt,
+            "checkpoint_interval_steps": int(max(self.checkpoint_interval_steps, 1)),
+        }
+        self.manifest_path.write_text(json.dumps(_json_safe(manifest), indent=2), encoding="utf-8")
+
+    def artifact_paths(self) -> dict[str, Any]:
+        return {
+            "logs_dir": str(self.logs_dir),
+            "run_manifest_path": str(self.manifest_path),
+            "events_path": str(self.events_path),
+            "latest_checkpoint_path": str(self.latest_checkpoint_path),
+            "checkpoint_paths": [str(path) for path in self.checkpoint_paths],
+            "run_summary_path": str(self.run_summary_path),
+        }
+
+    def log_event(self, event: dict[str, Any]) -> None:
+        dynamic_event = {
+            "run_id": self.run_id,
+            "timestamp": _utc_now_iso(),
+            "phase": event.get("phase"),
+            "step": event.get("step"),
+            "train_episode_index": event.get("train_episode_index"),
+            "problem_id": event.get("problem_id"),
+            "problem_family": event.get("problem_family"),
+            "difficulty": event.get("difficulty"),
+            "teacher_prompt": event.get("teacher_prompt"),
+            "solver_completion": event.get("solver_completion"),
+            "extracted_code": event.get("extracted_code"),
+            "reward": event.get("reward"),
+            "pass_rate": event.get("pass_rate"),
+            "visible_pass_rate": event.get("visible_pass_rate"),
+            "execution_status": event.get("execution_status"),
+            "efficiency_score": event.get("efficiency_score"),
+            "optimization_hints": event.get("optimization_hints", []),
+            "feedback": event.get("feedback"),
+        }
+        with self.events_path.open("a", encoding="utf-8") as handle:
+            handle.write(json.dumps(_json_safe(dynamic_event)) + "\n")
+
+        self._latest_event = dynamic_event
+        if dynamic_event.get("reward") is not None:
+            self._recent_rewards.append(float(dynamic_event["reward"]))
+        if dynamic_event.get("pass_rate") is not None:
+            self._recent_pass_rates.append(float(dynamic_event["pass_rate"]))
+        if dynamic_event.get("efficiency_score") is not None:
+            self._recent_efficiency_scores.append(float(dynamic_event["efficiency_score"]))
+
+    def record_progress(self, progress: dict[str, Any]) -> None:
+        self._latest_progress.update({key: value for key, value in progress.items() if value is not None})
+        completed_steps = int(self._latest_progress.get("completed_steps", 0) or 0)
+        interval = int(max(self.checkpoint_interval_steps, 1))
+        if completed_steps > 0 and completed_steps % interval == 0 and completed_steps != self._last_checkpoint_step:
+            self._write_checkpoint(completed_steps)
+
+    def finalize(self, *, reward_curve_csv: Path | None = None, final_metrics: dict[str, Any] | None = None) -> Path:
+        copied_reward_curve = None
+        if reward_curve_csv is not None and reward_curve_csv.exists():
+            copied_reward_curve = self.logs_dir / "reward_curve.csv"
+            if reward_curve_csv.resolve() != copied_reward_curve.resolve():
+                shutil.copy2(reward_curve_csv, copied_reward_curve)
+            else:
+                copied_reward_curve = reward_curve_csv
+
+        summary = {
+            "run_id": self.run_id,
+            "finished_at": _utc_now_iso(),
+            "artifact_paths": self.artifact_paths(),
+            "reward_curve_csv": str(copied_reward_curve) if copied_reward_curve else None,
+            "latest_progress": self._latest_progress,
+            "latest_event": self._latest_event,
+            "rolling_metrics": self._rolling_metrics(),
+            "final_metrics": final_metrics or {},
+        }
+        self.run_summary_path.write_text(json.dumps(_json_safe(summary), indent=2), encoding="utf-8")
+        return self.run_summary_path
+
+    def _write_checkpoint(self, step: int) -> None:
+        checkpoint_payload = {
+            "run_id": self.run_id,
+            "timestamp": _utc_now_iso(),
+            "step": int(step),
+            "phase": self._latest_progress.get("phase"),
+            "total_steps": self._latest_progress.get("total_steps"),
+            "remaining_steps": self._latest_progress.get("remaining_steps"),
+            "progress_ratio": self._latest_progress.get("progress_ratio"),
+            "current_epoch": self._latest_progress.get("current_epoch"),
+            "current_difficulty": self._latest_progress.get("current_difficulty"),
+            "curriculum_level": self._latest_progress.get("curriculum_level"),
+            "train_episode_index": self._latest_progress.get("train_episode_index"),
+            "last_problem_id": self._latest_progress.get("last_problem_id"),
+            "last_problem_family": self._latest_progress.get("last_problem_family"),
+            "last_execution_status": self._latest_progress.get("last_execution_status"),
+            "rolling_metrics": self._rolling_metrics(),
+            "artifact_paths": {
+                "events_path": str(self.events_path),
+                "latest_checkpoint_path": str(self.latest_checkpoint_path),
+            },
+        }
+        checkpoint_path = self.logs_dir / f"checkpoint_step_{step:05d}.json"
+        checkpoint_path.write_text(json.dumps(_json_safe(checkpoint_payload), indent=2), encoding="utf-8")
+        self.latest_checkpoint_path.write_text(json.dumps(_json_safe(checkpoint_payload), indent=2), encoding="utf-8")
+        self.checkpoint_paths.append(checkpoint_path)
+        self._last_checkpoint_step = step
+
+    def _rolling_metrics(self) -> dict[str, Any]:
+        def _average(values: deque[float]) -> float | None:
+            if not values:
+                return None
+            return round(sum(values) / len(values), 4)
+
+        return {
+            "avg_reward": _average(self._recent_rewards),
+            "avg_pass_rate": _average(self._recent_pass_rates),
+            "avg_efficiency_score": _average(self._recent_efficiency_scores),
+        }

training/train_grpo.py

@@ -11,6 +11,7 @@ from typing import Any, Callable
 from env.adapt_env import AdaptEnvironment, MAX_STEPS_PER_EPISODE
 from env.generator import DIFFICULTY_LABELS, GeneratorAgent
 from models import AdaptAction
+from training.trace_logging import TraceArtifactLogger
 
 SYSTEM_PROMPT = """You are the Solver Agent for ADAPT.
 Write only runnable Python code.
@@ -44,6 +45,8 @@ class TrainingConfig:
     wandb_run_name: str | None = None
     generator_mode: str = "reward_aware"
     non_deterministic_generator: bool = False
+    trace_logging_enabled: bool = True
+    checkpoint_log_interval_steps: int = 10
 
     def to_dict(self) -> dict[str, Any]:
         return asdict(self)
@@ -60,6 +63,7 @@ TRAINING_PRESETS: dict[str, dict[str, Any]] = {
         "baseline_eval": False,
         "disable_wandb": True,
         "output_dir": "outputs_smoke",
+        "checkpoint_log_interval_steps": 2,
     },
     "default": {},
 }
@@ -156,6 +160,8 @@ def namespace_to_config(args: argparse.Namespace) -> TrainingConfig:
         wandb_run_name=args.wandb_run_name,
         generator_mode=args.generator_mode,
         non_deterministic_generator=args.non_deterministic_generator,
+        trace_logging_enabled=args.trace_logging_enabled,
+        checkpoint_log_interval_steps=args.checkpoint_log_interval_steps,
     )
 
 
@@ -329,12 +335,27 @@ class TrainingLogger:
     use_wandb: bool = True
     wandb_project: str = "adapt-dsa-tutor"
     wandb_run_name: str | None = None
+    run_id: str | None = None
+    training_config: dict[str, Any] = field(default_factory=dict)
+    model_identifiers: dict[str, Any] = field(default_factory=dict)
+    trace_logging_enabled: bool = True
+    checkpoint_log_interval_steps: int = 10
     rows: list[dict[str, Any]] = field(default_factory=list)
     global_step: int = 0
     _wandb_run: Any = field(default=None, init=False, repr=False)
+    _trace_logger: TraceArtifactLogger | None = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
         self.output_dir.mkdir(parents=True, exist_ok=True)
+        if self.trace_logging_enabled and self.run_id:
+            self._trace_logger = TraceArtifactLogger(
+                run_id=self.run_id,
+                output_dir=self.output_dir,
+                training_config=dict(self.training_config),
+                model_identifiers=dict(self.model_identifiers),
+                system_prompt=SYSTEM_PROMPT,
+                checkpoint_interval_steps=int(max(self.checkpoint_log_interval_steps, 1)),
+            )
         if not self.use_wandb:
             return
         try:
@@ -381,10 +402,37 @@ class TrainingLogger:
         if extra:
             row.update(extra)
         self.rows.append(row)
+        if self._trace_logger is not None:
+            self._trace_logger.log_event(
+                {
+                    "phase": phase,
+                    "step": self.global_step,
+                    "train_episode_index": extra.get("train_episode_index") if extra else None,
+                    "problem_id": row.get("problem_id"),
+                    "problem_family": row.get("problem_family"),
+                    "difficulty": row.get("difficulty_tier"),
+                    "teacher_prompt": row.get("teacher_prompt"),
+                    "solver_completion": row.get("solver_completion"),
+                    "extracted_code": row.get("extracted_code"),
+                    "reward": row.get("episode_reward"),
+                    "pass_rate": row.get("pass_rate"),
+                    "visible_pass_rate": row.get("visible_pass_rate"),
+                    "execution_status": row.get("execution_status"),
+                    "efficiency_score": row.get("efficiency_score"),
+                    "optimization_hints": row.get("optimization_hints", []),
+                    "feedback": row.get("feedback"),
+                }
+            )
         if self._wandb_run is not None:
             self._wandb_run.log(row, step=self.global_step)
         self.global_step += 1
 
+    def record_progress(self, updates: dict[str, Any]) -> dict[str, Any]:
+        if self._trace_logger is None:
+            return {}
+        self._trace_logger.record_progress(updates)
+        return self._trace_logger.artifact_paths()
+
     def write_csv(self) -> Path:
         output_path = self.output_dir / "reward_curve.csv"
         fieldnames: list[str] = []
@@ -402,11 +450,37 @@ class TrainingLogger:
         if self._wandb_run is not None:
             self._wandb_run.finish()
 
+    def finalize_trace_artifacts(
+        self,
+        *,
+        reward_curve_csv: Path | None = None,
+        final_metrics: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        if self._trace_logger is None:
+            return {}
+        self._trace_logger.finalize(reward_curve_csv=reward_curve_csv, final_metrics=final_metrics)
+        return self._trace_logger.artifact_paths()
+
 
 def build_dataset(size: int, controller: GeneratorController, curriculum: CurriculumManager) -> GeneratorRolloutDataset:
     return GeneratorRolloutDataset(size=size, controller=controller, curriculum=curriculum)
 
 
+def extract_optimization_hints(feedback: str) -> list[str]:
+    lines = [line.strip() for line in feedback.splitlines()]
+    hints: list[str] = []
+    capture = False
+    for line in lines:
+        if line == "Optimization hints:":
+            capture = True
+            continue
+        if capture and line.startswith("- "):
+            hints.append(line[2:])
+        elif capture and line:
+            break
+    return hints
+
+
 def build_reward_func(
     curriculum: CurriculumManager,
     controller: GeneratorController,
@@ -453,6 +527,13 @@ def build_reward_func(
                 extra={
                     "generator_reward": round(float(observation.generator_reward_signal), 4),
                     "problem_id": problem["problem_id"],
+                    "teacher_prompt": prompt,
+                    "solver_completion": completion,
+                    "extracted_code": extract_code(completion),
+                    "feedback": observation.feedback,
+                    "efficiency_score": observation.reward_components.get("efficiency_score"),
+                    "optimization_hints": extract_optimization_hints(observation.feedback),
+                    "train_episode_index": int(controller.history["episode_index"]),
                 },
             )
             if progress_callback is not None:
@@ -552,6 +633,9 @@ def run_policy_evaluation(
             session_id=env.session_id,
             generator_mode=generator_mode,
         )
+        last_prompt = ""
+        last_completion = ""
+        last_code = ""
 
         for _ in range(MAX_STEPS_PER_EPISODE):
             prompt = build_solver_prompt(observation.model_dump())
@@ -561,10 +645,13 @@ def run_policy_evaluation(
                 prompt=prompt,
                 max_new_tokens=max_new_tokens,
             )
+            last_prompt = prompt
+            last_completion = completion
+            last_code = extract_code(completion)
             observation = env.step(
                 AdaptAction(
                     session_id=env.session_id,
-                    code=extract_code(completion),
+                    code=last_code,
                 )
             )
             if observation.done:
@@ -591,6 +678,13 @@ def run_policy_evaluation(
             extra={
                 "generator_reward": round(float(observation.generator_reward_signal), 4),
                 "problem_id": problem["problem_id"],
+                "teacher_prompt": last_prompt,
+                "solver_completion": last_completion,
+                "extracted_code": last_code,
+                "feedback": observation.feedback,
+                "efficiency_score": observation.reward_components.get("efficiency_score"),
+                "optimization_hints": extract_optimization_hints(observation.feedback),
+                "train_episode_index": int(controller.history["episode_index"]),
             },
         )
 
@@ -615,11 +709,17 @@ def print_evaluation_summary(baseline: dict[str, Any], trained: dict[str, Any])
 def run_training(
     config: TrainingConfig | argparse.Namespace,
     *,
+    run_id: str | None = None,
     progress_callback: Callable[[dict[str, Any]], None] | None = None,
 ) -> dict[str, Any]:
     if isinstance(config, argparse.Namespace):
         config = namespace_to_config(config)
 
+    try:
+        import torch
+    except ImportError as exc:
+        raise RuntimeError("Training requires `torch` to be installed.") from exc
+
     try:
         from trl import GRPOConfig, GRPOTrainer
         from unsloth import FastLanguageModel, PatchFastRL
@@ -634,13 +734,20 @@ def run_training(
 
     PatchFastRL("GRPO", FastLanguageModel)
 
+    use_cuda = torch.cuda.is_available()
+    use_bf16 = bool(config.bf16) or (use_cuda and hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported())
+    model_dtype = torch.bfloat16 if use_bf16 else (torch.float16 if use_cuda else torch.float32)
+
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name=config.model_name,
         max_seq_length=config.max_seq_length,
+        dtype=model_dtype,
         load_in_4bit=not config.disable_4bit,
     )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
+    if hasattr(model, "config"):
+        model.config.torch_dtype = model_dtype
 
     model = FastLanguageModel.get_peft_model(
         model,
@@ -661,22 +768,36 @@ def run_training(
         use_wandb=not config.disable_wandb,
         wandb_project=config.wandb_project,
         wandb_run_name=config.wandb_run_name,
+        run_id=run_id,
+        training_config=config.to_dict(),
+        model_identifiers={
+            "model_name": config.model_name,
+            "generator_mode": config.generator_mode,
+        },
+        trace_logging_enabled=config.trace_logging_enabled,
+        checkpoint_log_interval_steps=config.checkpoint_log_interval_steps,
     )
 
+    def emit_progress(update: dict[str, Any]) -> None:
+        artifact_paths = logger.record_progress(update)
+        if progress_callback is not None:
+            payload = dict(update)
+            payload.update(artifact_paths)
+            progress_callback(payload)
+
     baseline_summary = {"easy": 0.0, "medium": 0.0, "hard": 0.0, "overall": 0.0}
     trained_summary = {"easy": 0.0, "medium": 0.0, "hard": 0.0, "overall": 0.0}
 
     if config.baseline_eval:
         FastLanguageModel.for_inference(model)
-        if progress_callback is not None:
-            progress_callback(
-                {
-                    "phase": "baseline_eval",
-                    "status": "running",
-                    "completed_steps": 0,
-                    "total_steps": int(config.max_steps),
-                }
-            )
+        emit_progress(
+            {
+                "phase": "baseline_eval",
+                "status": "running",
+                "completed_steps": 0,
+                "total_steps": int(config.max_steps),
+            }
+        )
         baseline_summary = run_policy_evaluation(
             model=model,
             tokenizer=tokenizer,
@@ -688,14 +809,13 @@ def run_training(
             max_new_tokens=config.eval_max_new_tokens,
         )
         print(f"[baseline_eval] {json.dumps(baseline_summary)}")
-        if progress_callback is not None:
-            progress_callback(
-                {
-                    "phase": "baseline_eval",
-                    "status": "completed",
-                    "baseline_summary": baseline_summary,
-                }
-            )
+        emit_progress(
+            {
+                "phase": "baseline_eval",
+                "status": "completed",
+                "baseline_summary": baseline_summary,
+            }
+        )
 
     training_args = GRPOConfig(
         output_dir=str(output_dir),
@@ -707,49 +827,47 @@ def run_training(
         max_completion_length=config.max_completion_length,
         max_steps=config.max_steps,
         logging_steps=1,
-        bf16=config.bf16,
+        bf16=use_bf16,
+        fp16=use_cuda and not use_bf16,
         report_to=[],
     )
 
     class ProgressCallback(TrainerCallback):
         def on_train_begin(self, args, state, control, **kwargs):  # type: ignore[override]
             del args, control, kwargs
-            if progress_callback is not None:
-                progress_callback(
-                    {
-                        "phase": "train",
-                        "status": "running",
-                        "completed_steps": int(getattr(state, "global_step", 0)),
-                        "total_steps": int(config.max_steps),
-                        "current_epoch": float(getattr(state, "epoch", 0.0) or 0.0),
-                    }
-                )
+            emit_progress(
+                {
+                    "phase": "train",
+                    "status": "running",
+                    "completed_steps": int(getattr(state, "global_step", 0)),
+                    "total_steps": int(config.max_steps),
+                    "current_epoch": float(getattr(state, "epoch", 0.0) or 0.0),
+                }
+            )
 
         def on_step_end(self, args, state, control, **kwargs):  # type: ignore[override]
             del args, control, kwargs
-            if progress_callback is not None:
-                progress_callback(
-                    {
-                        "phase": "train",
-                        "status": "running",
-                        "completed_steps": int(getattr(state, "global_step", 0)),
-                        "total_steps": int(config.max_steps),
-                        "current_epoch": float(getattr(state, "epoch", 0.0) or 0.0),
-                    }
-                )
+            emit_progress(
+                {
+                    "phase": "train",
+                    "status": "running",
+                    "completed_steps": int(getattr(state, "global_step", 0)),
+                    "total_steps": int(config.max_steps),
+                    "current_epoch": float(getattr(state, "epoch", 0.0) or 0.0),
+                }
+            )
 
         def on_train_end(self, args, state, control, **kwargs):  # type: ignore[override]
             del args, control, kwargs
-            if progress_callback is not None:
-                progress_callback(
-                    {
-                        "phase": "train",
-                        "status": "completed",
-                        "completed_steps": int(getattr(state, "global_step", 0)),
-                        "total_steps": int(config.max_steps),
-                        "current_epoch": float(getattr(state, "epoch", 0.0) or 0.0),
-                    }
-                )
+            emit_progress(
+                {
+                    "phase": "train",
+                    "status": "completed",
+                    "completed_steps": int(getattr(state, "global_step", 0)),
+                    "total_steps": int(config.max_steps),
+                    "current_epoch": float(getattr(state, "epoch", 0.0) or 0.0),
+                }
+            )
 
     trainer = GRPOTrainer(
         model=model,
@@ -765,15 +883,14 @@ def run_training(
 
     if config.baseline_eval:
         FastLanguageModel.for_inference(model)
-        if progress_callback is not None:
-            progress_callback(
-                {
-                    "phase": "trained_eval",
-                    "status": "running",
-                    "completed_steps": int(config.max_steps),
-                    "total_steps": int(config.max_steps),
-                }
-            )
+        emit_progress(
+            {
+                "phase": "trained_eval",
+                "status": "running",
+                "completed_steps": int(config.max_steps),
+                "total_steps": int(config.max_steps),
+            }
+        )
         trained_summary = run_policy_evaluation(
             model=model,
             tokenizer=tokenizer,
@@ -786,16 +903,23 @@ def run_training(
         )
         print(f"[trained_eval] {json.dumps(trained_summary)}")
         print_evaluation_summary(baseline_summary, trained_summary)
-        if progress_callback is not None:
-            progress_callback(
-                {
-                    "phase": "trained_eval",
-                    "status": "completed",
-                    "trained_summary": trained_summary,
-                }
-            )
+        emit_progress(
+            {
+                "phase": "trained_eval",
+                "status": "completed",
+                "trained_summary": trained_summary,
+            }
+        )
 
     csv_path = logger.write_csv()
+    trace_artifact_paths = logger.finalize_trace_artifacts(
+        reward_curve_csv=csv_path,
+        final_metrics={
+            "baseline_summary": baseline_summary,
+            "trained_summary": trained_summary,
+            "completed_steps": int(config.max_steps),
+        },
+    )
     logger.close()
     print(f"[artifacts] reward curve CSV written to {csv_path}")
 
@@ -803,6 +927,7 @@ def run_training(
         "config": config.to_dict(),
         "output_dir": str(output_dir.resolve()),
         "reward_curve_csv": str(csv_path.resolve()),
+        **trace_artifact_paths,
         "baseline_summary": baseline_summary,
         "trained_summary": trained_summary,
         "completed_steps": int(config.max_steps),
@@ -832,6 +957,8 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--disable-wandb", action="store_true")
     parser.add_argument("--wandb-project", default="adapt-dsa-tutor")
     parser.add_argument("--wandb-run-name", default=None)
+    parser.add_argument("--trace-logging-enabled", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--checkpoint-log-interval-steps", type=int, default=10)
     parser.add_argument(
         "--generator-mode",
         choices=["heuristic", "reward_aware"],