feat: enhance CyberSecurity_OWASP observation model with scenario prompt, improve GRPO batch configuration validation, and add scenario grouping for adaptive difficulty curriculum

models.py

@@ -38,6 +38,7 @@ class CyberSecurityOWASPObservation(Observation):
     phase: CyberSecurityOWASPPhase = "discover"
     message: str = ""
     task_brief: str = ""
+    scenario_prompt: str = ""
     visible_policy_hint: dict[str, Any] = Field(default_factory=dict)
     workspace_summary: dict[str, Any] = Field(default_factory=dict)
     available_actions: list[str] = Field(default_factory=list)

scripts/modal_train_grpo.py

@@ -71,6 +71,55 @@ def _hf_model_cache_path(model_name: str) -> pathlib.Path:
     return HF_HUB_CACHE_DIR / f"models--{model_name.replace('/', '--')}"
 
 
+def _resolve_grpo_batch_config(
+    *,
+    per_device_train_batch_size: int,
+    gradient_accumulation_steps: int,
+    num_generations: int,
+    world_size: int = 1,
+) -> tuple[int, int]:
+    if num_generations < 1:
+        raise ValueError("--num-generations must be at least 1.")
+    if per_device_train_batch_size < 1:
+        raise ValueError("--per-device-train-batch-size must be at least 1.")
+    if world_size < 1:
+        raise ValueError("world_size must be at least 1.")
+
+    resolved_gradient_accumulation_steps = (
+        gradient_accumulation_steps
+        if gradient_accumulation_steps > 0
+        else max(2, num_generations)
+    )
+    if resolved_gradient_accumulation_steps < 1:
+        raise ValueError("--gradient-accumulation-steps must be at least 1.")
+
+    effective_batch_size = (
+        per_device_train_batch_size
+        * resolved_gradient_accumulation_steps
+        * world_size
+    )
+    if effective_batch_size % num_generations:
+        raise ValueError(
+            "Invalid GRPO batch shape: "
+            "per_device_train_batch_size * gradient_accumulation_steps * world_size "
+            f"must be divisible by num_generations. Got "
+            f"{per_device_train_batch_size} * "
+            f"{resolved_gradient_accumulation_steps} * {world_size} = "
+            f"{effective_batch_size}, which is not divisible by {num_generations}."
+        )
+    return resolved_gradient_accumulation_steps, effective_batch_size
+
+
+def _validate_vllm_config(*, use_vllm: bool, vllm_gpu_memory_utilization: float) -> None:
+    if not use_vllm:
+        return
+    if not 0.0 < vllm_gpu_memory_utilization <= 0.95:
+        raise ValueError(
+            "--vllm-gpu-memory-utilization must be in the interval (0.0, 0.95] "
+            "when --use-vllm is enabled."
+        )
+
+
 def _configure_modal_cache_env() -> dict[str, str]:
     values = {
         "HF_HOME": str(HF_HOME_DIR),
@@ -374,22 +423,20 @@ def verify_modal_scenario_cache_for_training(
     resolved_difficulty = int(scenario_profile["difficulty"])
     cache = ScenarioCache(SCENARIO_CACHE_DIR, settings=settings)
     coverage = cache.assert_coverage(split=split, difficulty=resolved_difficulty)
-    available_scenarios = int(
-        coverage.get("counts", {})
-        .get(split, {})
-        .get(str(resolved_difficulty), 0)
-    )
-    if available_scenarios < dataset_size:
-        raise RuntimeError(
-            "Scenario cache does not cover this Modal dataset. Run "
-            "--mode prepare-cache with a larger per-bucket count before training. "
-            f"available={available_scenarios}, requested_dataset_size={dataset_size}, "
-            f"split={split}, difficulty={resolved_difficulty}"
-        )
+    entries = cache.validated_entries(split=split, difficulty=resolved_difficulty)
+    if not entries:
+        entries = cache.validated_entries(split=split)
+    if not entries:
+        raise RuntimeError(f"No validated scenario cache entries found for split={split!r}.")
+    sample_entry = entries[0]
 
     env = CybersecurityOwaspEnvironment()
     try:
-        obs = env.reset(seed=seed_start, split=split, difficulty=difficulty)
+        obs = env.reset(
+            seed=int(sample_entry["seed"]),
+            split=str(sample_entry["split"]),
+            difficulty=int(sample_entry["difficulty"]),
+        )
         if not env.state.cache_hit:
             raise RuntimeError("Scenario cache preflight reset did not hit cache.")
         if env.state.metrics.get("scenario_compile_latency_ms", 0.0):
@@ -413,9 +460,10 @@ def verify_modal_scenario_cache_for_training(
         "scenario_cache_dir": str(SCENARIO_CACHE_DIR),
         "scenario_cache_mode": "require",
         "split": split,
-        "difficulty": resolved_difficulty,
+        "difficulty": "adaptive",
+        "initial_difficulty": resolved_difficulty,
         "dataset_size": dataset_size,
-        "available_scenarios": available_scenarios,
+        "available_scenarios": len(cache.validated_entries(split=split)),
         "coverage": coverage,
         "sample_reset": sample,
     }
@@ -480,6 +528,11 @@ def train_cybersecurity_owasp_grpo(
     trackio_space_id: str = "Humanlearning/CyberSecurity_OWASP-trackio",
     trackio_project: str = "CyberSecurity_OWASP-grpo",
     num_generations: int = 6,
+    per_device_train_batch_size: int = 1,
+    gradient_accumulation_steps: int = 0,
+    use_vllm: bool = False,
+    vllm_gpu_memory_utilization: float = 0.2,
+    trace_log_every: int = 5,
     seed_start: int = 0,
     git_sha: str = "nogit",
     run_name: str = "",
@@ -495,6 +548,21 @@ def train_cybersecurity_owasp_grpo(
 
     model_name = _ensure_gemma4_model(model_name)
     cache_env = _configure_modal_cache_env()
+    world_size = int(os.environ.get("WORLD_SIZE", "1") or "1")
+    (
+        resolved_gradient_accumulation_steps,
+        effective_train_batch_size,
+    ) = _resolve_grpo_batch_config(
+        per_device_train_batch_size=per_device_train_batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        num_generations=num_generations,
+        world_size=world_size,
+    )
+    _validate_vllm_config(
+        use_vllm=use_vllm,
+        vllm_gpu_memory_utilization=vllm_gpu_memory_utilization,
+    )
+    trace_log_every = max(0, int(trace_log_every))
 
     import torch
     from unsloth import FastVisionModel
@@ -524,6 +592,10 @@ def train_cybersecurity_owasp_grpo(
         log_trackio_metrics,
         train_metric_aliases,
     )
+    from training.grpo_curriculum import (
+        ScenarioGroupRegistry,
+        build_scenario_group_rows,
+    )
 
     transformers_hub.TRANSFORMERS_CACHE = cache_env["HF_HUB_CACHE"]
 
@@ -585,18 +657,14 @@ def train_cybersecurity_owasp_grpo(
         split=split,
         difficulty=int(scenario_profile["difficulty"]),
     )
-    available_scenarios = int(
-        scenario_cache_coverage.get("counts", {})
-        .get(split, {})
-        .get(str(int(scenario_profile["difficulty"])), 0)
+    scenario_entries = scenario_cache.validated_entries(split=split)
+    scenario_registry = ScenarioGroupRegistry(
+        scenario_entries,
+        split=split,
+        initial_difficulty=int(scenario_profile["difficulty"]),
+        rng_seed=seed_start,
+        max_level=scenario_settings.curriculum.difficulty_bucket_count - 1,
     )
-    if available_scenarios < dataset_size:
-        raise RuntimeError(
-            "Scenario cache does not cover this Modal dataset. Run "
-            "--mode prepare-cache with a larger per-bucket count before training. "
-            f"available={available_scenarios}, requested_dataset_size={dataset_size}, "
-            f"split={split}, difficulty={scenario_profile['difficulty']}"
-        )
 
     training_prompt = (
         "You are a defensive AppSec repair agent in the local CyberSecurity_OWASP "
@@ -608,15 +676,14 @@ def train_cybersecurity_owasp_grpo(
     )
 
     dataset = Dataset.from_list(
-        [
-            {
-                "prompt": [{"role": "user", "content": training_prompt}],
-                "seed": seed_start + index,
-                "difficulty": difficulty,
-                "split": split,
-            }
-            for index in range(dataset_size)
-        ]
+        build_scenario_group_rows(
+            dataset_size=dataset_size,
+            training_prompt=training_prompt,
+            seed_start=seed_start,
+            split=split,
+            difficulty=difficulty,
+            difficulty_policy="adaptive",
+        )
     )
 
     def _state_snapshot(env: CybersecurityOwaspEnvironment) -> dict[str, Any]:
@@ -627,8 +694,10 @@ def train_cybersecurity_owasp_grpo(
             "seed": state.seed,
             "split": state.split,
             "difficulty": state.difficulty,
+            "difficulty_tier": state.difficulty_tier,
             "domain": state.domain,
             "bug_family": state.bug_family,
+            "template_id": state.template_id,
             "cache_hit": state.cache_hit,
             "scenario_hash": state.scenario_hash,
             "phase": state.phase,
@@ -647,18 +716,30 @@ def train_cybersecurity_owasp_grpo(
             self.done = False
             self.success = False
             self.invalid_actions = 0
+            self.scenario_group_id = -1
+            self.scenario_assignment: dict[str, Any] = {}
             self.trace_messages: list[dict[str, str]] = []
             self.trace_metadata: dict[str, Any] = {}
 
         def reset(self, **kwargs) -> str:
-            seed = int(kwargs.get("seed", seed_start))
-            current_difficulty = int(kwargs.get("difficulty", difficulty))
-            current_split = str(kwargs.get("split", split))
+            group_id = int(kwargs.get("scenario_group_id", kwargs.get("seed", seed_start)))
+            assignment = scenario_registry.assignment_for(
+                scenario_group_id=group_id,
+                requested_seed=int(kwargs.get("seed", seed_start)),
+                requested_difficulty=int(kwargs.get("difficulty", difficulty)),
+                split=str(kwargs.get("split", split)),
+                difficulty_policy=str(kwargs.get("difficulty_policy", "adaptive")),
+            )
+            seed = int(assignment["seed"])
+            current_difficulty = int(assignment["difficulty"])
+            current_split = str(assignment["split"])
             obs = self._env.reset(
                 seed=seed,
                 split=current_split,
                 difficulty=current_difficulty,
             )
+            self.scenario_group_id = group_id
+            self.scenario_assignment = assignment
             self.reward = 0.0
             self.reward_breakdown = {}
             self.done = bool(obs.done)
@@ -668,18 +749,21 @@ def train_cybersecurity_owasp_grpo(
                 {
                     "role": "user",
                     "content": (
-                        f"{training_prompt}\n\nInitial observation:\n"
-                        f"Phase: {obs.phase}\n"
-                        f"Task: {obs.task_brief}\n"
-                        f"Available actions: {obs.available_actions}\n"
-                        f"Workspace summary: {obs.workspace_summary}\n"
-                        f"Policy hint: {obs.visible_policy_hint}\n"
-                        f"Message: {obs.message}"
+                        f"{training_prompt}\n\n"
+                        f"{obs.scenario_prompt}\n\n"
+                        f"Initial message: {obs.message}"
                     ),
                 }
             ]
             self.trace_metadata = _state_snapshot(self._env)
-            return obs.message
+            self.trace_metadata.update(
+                {
+                    "scenario_group_id": self.scenario_group_id,
+                    "scenario_assignment": dict(self.scenario_assignment),
+                    "scenario_prompt_length": len(obs.scenario_prompt),
+                }
+            )
+            return obs.scenario_prompt
 
         def _step(self, tool_name: str, arguments: dict[str, Any] | None = None) -> str:
             if self.done:
@@ -714,6 +798,8 @@ def train_cybersecurity_owasp_grpo(
                     "invalid_actions": self.invalid_actions,
                     "scenario_cache_hit": self._env.state.cache_hit,
                     "scenario_hash": self._env.state.scenario_hash,
+                    "scenario_group_id": self.scenario_group_id,
+                    "scenario_assignment": dict(self.scenario_assignment),
                 }
             )
             return obs.message
@@ -938,11 +1024,58 @@ def train_cybersecurity_owasp_grpo(
             )
             episode_records.append(record)
 
+        group_successes: dict[int, list[float]] = {}
+        for env in environments:
+            group_id = int(getattr(env, "scenario_group_id", -1))
+            if group_id < 0:
+                continue
+            group_successes.setdefault(group_id, []).append(1.0 if getattr(env, "success", False) else 0.0)
+        for group_id, successes in group_successes.items():
+            scenario_registry.record_group_outcome(group_id, _mean(successes))
+
+        batch_fingerprints = [
+            episode_trace_fingerprint(record)
+            for record in episode_records
+        ]
+        sampled_traces = []
+        seen_this_batch: set[str] = set()
+        duplicate_trace_suppressed_count = 0
+        for index, (env, record, reward, fingerprint) in enumerate(
+            zip(environments, episode_records, rewards, batch_fingerprints)
+        ):
+            if fingerprint in seen_this_batch or fingerprint in logged_trace_fingerprints:
+                duplicate_trace_suppressed_count += 1
+                continue
+            seen_this_batch.add(fingerprint)
+            if len(sampled_traces) < 4:
+                sampled_traces.append((index, env, record, reward, fingerprint))
+
+        should_log_trace_artifacts = trace_log_every > 0 and (
+            trace_step["value"] == 1
+            or trace_step["value"] % trace_log_every == 0
+        )
         canonical_metrics = aggregate_episode_metrics(episode_records)
         metrics = {
             **canonical_metrics,
             **train_metric_aliases(canonical_metrics),
+            **scenario_registry.metrics(
+                episode_records,
+                unique_trace_count=len(set(batch_fingerprints)),
+                duplicate_trace_suppressed_count=duplicate_trace_suppressed_count,
+            ),
         }
+        metrics["train/per_device_train_batch_size"] = float(per_device_train_batch_size)
+        metrics["train/gradient_accumulation_steps"] = float(
+            resolved_gradient_accumulation_steps
+        )
+        metrics["train/effective_train_batch_size"] = float(effective_train_batch_size)
+        metrics["train/num_generations"] = float(num_generations)
+        metrics["train/use_vllm"] = float(bool(use_vllm))
+        metrics["train/vllm_gpu_memory_utilization"] = (
+            float(vllm_gpu_memory_utilization) if use_vllm else 0.0
+        )
+        metrics["train/trace_log_every"] = float(trace_log_every)
+        metrics["train/trace_artifacts_logged"] = float(should_log_trace_artifacts)
         if rewards:
             metrics["train/reward_mean"] = _mean(rewards)
             metrics["train/reward_std"] = statistics.pstdev(rewards) if len(rewards) > 1 else 0.0
@@ -952,60 +1085,57 @@ def train_cybersecurity_owasp_grpo(
         except Exception as exc:
             print(f"Trackio metric logging skipped: {exc!r}")
 
-        sampled_traces = []
-        seen_this_batch: set[str] = set()
-        for index, (env, record, reward) in enumerate(zip(environments, episode_records, rewards)):
-            fingerprint = episode_trace_fingerprint(record)
-            if fingerprint in seen_this_batch or fingerprint in logged_trace_fingerprints:
-                continue
-            seen_this_batch.add(fingerprint)
-            logged_trace_fingerprints.add(fingerprint)
-            sampled_traces.append((index, env, record, reward, fingerprint))
-            if len(sampled_traces) >= 4:
-                break
-
-        try:
-            log_trace_table(
-                [record for _, _, record, _, _ in sampled_traces],
-                table_name="sample_traces",
-                step=trace_step["value"],
-            )
-        except Exception as exc:
-            print(f"Trackio sample trace table logging skipped: {exc!r}")
-
-        for index, env, _record, reward, fingerprint in sampled_traces:
-            messages = list(getattr(env, "trace_messages", []))
-            if index < len(completions):
-                completion_text = _completion_to_text(completions[index])
-                if completion_text:
-                    messages.append(
-                        {
-                            "role": "assistant",
-                            "content": f"Raw generated completion:\n{completion_text}",
-                        }
-                    )
-            metadata = dict(getattr(env, "trace_metadata", {}))
-            metadata.update(
-                {
-                    "sample_index": index,
-                    "reward": reward,
-                    "trace_step": trace_step["value"],
-                    "trace_fingerprint": fingerprint,
-                    "run_name": run_name,
-                }
-            )
+        if should_log_trace_artifacts and sampled_traces:
             try:
-                trackio.log(
-                    {
-                        f"cybersecurity_owasp_trace/sample_{index}": trackio.Trace(
-                            messages=messages,
-                            metadata=metadata,
-                        )
-                    },
+                log_trace_table(
+                    [record for _, _, record, _, _ in sampled_traces],
+                    table_name="sample_traces",
                     step=trace_step["value"],
                 )
             except Exception as exc:
-                print(f"Trackio trace logging skipped: {exc!r}")
+                print(f"Trackio sample trace table logging skipped: {exc!r}")
+
+            for index, env, _record, reward, fingerprint in sampled_traces:
+                logged_trace_fingerprints.add(fingerprint)
+                messages = list(getattr(env, "trace_messages", []))
+                if index < len(completions):
+                    completion_text = _completion_to_text(completions[index])
+                    if completion_text:
+                        messages.append(
+                            {
+                                "role": "assistant",
+                                "content": f"Raw generated completion:\n{completion_text}",
+                            }
+                        )
+                metadata = dict(getattr(env, "trace_metadata", {}))
+                metadata.update(
+                    {
+                        "sample_index": index,
+                        "reward": reward,
+                        "trace_step": trace_step["value"],
+                        "trace_fingerprint": fingerprint,
+                        "num_generations": num_generations,
+                        "run_name": run_name,
+                    }
+                )
+                try:
+                    trackio.log(
+                        {
+                            f"cybersecurity_owasp_trace/sample_{index}": trackio.Trace(
+                                messages=messages,
+                                metadata=metadata,
+                            )
+                        },
+                        step=trace_step["value"],
+                    )
+                except Exception as exc:
+                    print(f"Trackio trace logging skipped: {exc!r}")
+        elif sampled_traces:
+            print(
+                "Trackio trace artifacts throttled at reward callback "
+                f"{trace_step['value']}; set --trace-log-every 1 for every callback "
+                "or 0 to disable trace artifacts."
+            )
 
         if rewards:
             print(
@@ -1080,6 +1210,20 @@ def train_cybersecurity_owasp_grpo(
     print(f"Unsloth cache: {cache_env['UNSLOTH_CACHE_DIR']}")
     print(f"Triton cache: {cache_env['TRITON_CACHE_DIR']}")
     print(f"Hub push enabled: {push_to_hub}")
+    print(
+        "GRPO throughput config: "
+        f"per_device_train_batch_size={per_device_train_batch_size}, "
+        f"gradient_accumulation_steps={resolved_gradient_accumulation_steps}, "
+        f"num_generations={num_generations}, "
+        f"world_size={world_size}, "
+        f"effective_train_batch_size={effective_train_batch_size}"
+    )
+    print(
+        "Generation acceleration config: "
+        f"use_vllm={use_vllm}, "
+        f"vllm_gpu_memory_utilization={vllm_gpu_memory_utilization}, "
+        f"trace_log_every={trace_log_every}"
+    )
 
     expected_model_cache = _hf_model_cache_path(model_name)
     cache_hit = expected_model_cache.exists()
@@ -1109,13 +1253,36 @@ def train_cybersecurity_owasp_grpo(
 
     print(f"Loading model with Unsloth from_pretrained: {model_name}")
     model_api = FastVisionModel
+    model_load_values = {
+        "model_name": model_name,
+        "max_seq_length": max_seq_length,
+        "load_in_4bit": False,
+        "fast_inference": use_vllm,
+        "gpu_memory_utilization": vllm_gpu_memory_utilization if use_vllm else None,
+        "cache_dir": str(HF_HUB_CACHE_DIR),
+        "token": hf_token,
+    }
+    from_pretrained_parameters = inspect.signature(model_api.from_pretrained).parameters
+    from_pretrained_accepts_kwargs = any(
+        parameter.kind == inspect.Parameter.VAR_KEYWORD
+        for parameter in from_pretrained_parameters.values()
+    )
+    skipped_model_load_keys = sorted(
+        key
+        for key, value in model_load_values.items()
+        if value is not None
+        and key not in from_pretrained_parameters
+        and not from_pretrained_accepts_kwargs
+    )
+    if skipped_model_load_keys:
+        print(f"Skipping unsupported from_pretrained keys: {skipped_model_load_keys}")
     model, tokenizer = model_api.from_pretrained(
-        model_name=model_name,
-        max_seq_length=max_seq_length,
-        load_in_4bit=False,
-        fast_inference=False,
-        cache_dir=str(HF_HUB_CACHE_DIR),
-        token=hf_token,
+        **{
+            key: value
+            for key, value in model_load_values.items()
+            if value is not None
+            and (key in from_pretrained_parameters or from_pretrained_accepts_kwargs)
+        }
     )
     print("Model load complete.")
     cache_volume.commit()
@@ -1157,8 +1324,8 @@ def train_cybersecurity_owasp_grpo(
         "lr_scheduler_type": "linear",
         "optim": "adamw_8bit",
         "logging_steps": 1,
-        "per_device_train_batch_size": 1,
-        "gradient_accumulation_steps": max(2, num_generations),
+        "per_device_train_batch_size": per_device_train_batch_size,
+        "gradient_accumulation_steps": resolved_gradient_accumulation_steps,
         "num_generations": num_generations,
         "max_prompt_length": max_seq_length,
         "max_completion_length": max_completion_length,
@@ -1175,11 +1342,14 @@ def train_cybersecurity_owasp_grpo(
         "hub_strategy": "every_save",
         "gradient_checkpointing": True,
         "gradient_checkpointing_kwargs": {"use_reentrant": False},
+        "use_vllm": use_vllm,
+        "vllm_mode": "colocate",
+        "vllm_gpu_memory_utilization": vllm_gpu_memory_utilization,
         "epsilon": 0.2,
         "epsilon_high": 0.28,
         "delta": 1.5,
         "loss_type": "bnpo",
-        "mask_truncated_completions": True,
+        "mask_truncated_completions": False,
     }
     grpo_config_parameters = set(inspect.signature(GRPOConfig).parameters)
     skipped_config_keys = sorted(set(grpo_config_values) - grpo_config_parameters)
@@ -1269,6 +1439,12 @@ def train_cybersecurity_owasp_grpo(
         "model_name": model_name,
         "max_completion_length": max_completion_length,
         "num_generations": num_generations,
+        "per_device_train_batch_size": per_device_train_batch_size,
+        "gradient_accumulation_steps": resolved_gradient_accumulation_steps,
+        "effective_train_batch_size": effective_train_batch_size,
+        "use_vllm": int(bool(use_vllm)),
+        "vllm_gpu_memory_utilization": vllm_gpu_memory_utilization,
+        "trace_log_every": trace_log_every,
         "source_mode": source_mode,
         "repo_url": repo_url,
         "repo_branch": repo_branch,
@@ -1294,6 +1470,11 @@ def main(
     trackio_space_id: str = "Humanlearning/CyberSecurity_OWASP-trackio",
     trackio_project: str = "CyberSecurity_OWASP-grpo",
     num_generations: int = 6,
+    per_device_train_batch_size: int = 1,
+    gradient_accumulation_steps: int = 0,
+    use_vllm: bool = False,
+    vllm_gpu_memory_utilization: float = 0.2,
+    trace_log_every: int = 5,
     seed_start: int = 0,
     git_sha: str = "nogit",
     source_mode: str = "local",
@@ -1327,6 +1508,21 @@ def main(
     if mode != "train":
         raise ValueError("mode must be 'prepare-cache', 'train', or 'config'")
 
+    (
+        resolved_gradient_accumulation_steps,
+        effective_train_batch_size,
+    ) = _resolve_grpo_batch_config(
+        per_device_train_batch_size=per_device_train_batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        num_generations=num_generations,
+        world_size=1,
+    )
+    _validate_vllm_config(
+        use_vllm=use_vllm,
+        vllm_gpu_memory_utilization=vllm_gpu_memory_utilization,
+    )
+    trace_log_every = max(0, int(trace_log_every))
+
     trackio_space_id = trackio_space_id or os.environ.get(
         "TRACKIO_SPACE_ID",
         "Humanlearning/CyberSecurity_OWASP-trackio",
@@ -1392,6 +1588,19 @@ def main(
     print(f"Hub push enabled: {push_to_hub}")
     print(f"Model cache volume: {CACHE_VOLUME_NAME}")
     print(f"Scenario cache volume: {SCENARIO_CACHE_VOLUME_NAME}")
+    print(
+        "GRPO throughput config: "
+        f"per_device_train_batch_size={per_device_train_batch_size}, "
+        f"gradient_accumulation_steps={resolved_gradient_accumulation_steps}, "
+        f"num_generations={num_generations}, "
+        f"effective_train_batch_size={effective_train_batch_size}"
+    )
+    print(
+        "Generation acceleration config: "
+        f"use_vllm={use_vllm}, "
+        f"vllm_gpu_memory_utilization={vllm_gpu_memory_utilization}, "
+        f"trace_log_every={trace_log_every}"
+    )
     print("Launch phases:")
     print(
         "1. Modal image build/validation: happens before remote Python logs; "
@@ -1421,6 +1630,11 @@ def main(
         trackio_space_id=trackio_space_id,
         trackio_project=trackio_project,
         num_generations=num_generations,
+        per_device_train_batch_size=per_device_train_batch_size,
+        gradient_accumulation_steps=resolved_gradient_accumulation_steps,
+        use_vllm=use_vllm,
+        vllm_gpu_memory_utilization=vllm_gpu_memory_utilization,
+        trace_log_every=trace_log_every,
         seed_start=seed_start,
         git_sha=git_sha,
         run_name=run_name,

server/CyberSecurity_OWASP_environment.py

@@ -373,13 +373,15 @@ class CybersecurityOwaspEnvironment(
         visible_test_result: str | None = None,
         done_reason: str | None = None,
     ) -> CyberSecurityOWASPObservation:
+        available_actions = sorted(ALLOWED_TOOLS[self._state.phase])
         return CyberSecurityOWASPObservation(
             phase=self._state.phase,
             message=message,
             task_brief=self._task_brief,
+            scenario_prompt=self._scenario_prompt(available_actions),
             visible_policy_hint=self._visible_policy_hint,
             workspace_summary=self._workspace_summary,
-            available_actions=sorted(ALLOWED_TOOLS[self._state.phase]),
+            available_actions=available_actions,
             last_tool_result=message,
             last_action_valid=valid,
             last_action_error=error,
@@ -388,7 +390,60 @@ class CybersecurityOwaspEnvironment(
             done_reason=done_reason,
             done=self._state.done,
             reward=reward,
-            metadata={"episode_id": self._state.episode_id, "step_count": self._state.step_count},
+            metadata={
+                "episode_id": self._state.episode_id,
+                "step_count": self._state.step_count,
+                "seed": self._state.seed,
+                "split": self._state.split,
+                "difficulty": self._state.difficulty,
+                "difficulty_tier": self._state.difficulty_tier,
+                "domain": self._state.domain,
+                "bug_family": self._state.bug_family,
+                "template_id": self._state.template_id,
+                "scenario_hash": self._state.scenario_hash,
+            },
+        )
+
+    def _scenario_prompt(self, available_actions: list[str]) -> str:
+        users = self._visible_policy_hint.get("fixture_aliases", {}).get("users", {})
+        resources = self._visible_policy_hint.get("fixture_aliases", {}).get("resources", {})
+        visible_policy = {
+            "domain": self._visible_policy_hint.get("domain", self._state.domain),
+            "policy_rules": list(self._visible_policy_hint.get("policy_rules", [])),
+            "public_routes": list(self._visible_policy_hint.get("public_routes", [])),
+            "fixture_aliases": {
+                "users": sorted(str(key) for key in users),
+                "resources": sorted(str(key) for key in resources),
+            },
+        }
+        prompt = {
+            "environment": "CyberSecurity_OWASP",
+            "task": self._task_brief,
+            "scenario": {
+                "task_id": self._state.task_id,
+                "seed": self._state.seed,
+                "split": self._state.split,
+                "difficulty": self._state.difficulty,
+                "difficulty_tier": self._state.difficulty_tier,
+                "domain": self._state.domain,
+                "bug_family": self._state.bug_family,
+                "template_id": self._state.template_id,
+                "scenario_hash": self._state.scenario_hash,
+            },
+            "visible_policy_hint": visible_policy,
+            "workspace_summary": self._workspace_summary,
+            "available_actions": available_actions,
+            "instructions": [
+                "Use only the local generated application and the listed tools.",
+                "Discover the authorization defect with local evidence before patching.",
+                "Preserve legitimate owner/admin flows and intentionally public routes.",
+                "Submit exactly one secure, policy-aligned fix when ready.",
+            ],
+        }
+        return "CyberSecurity_OWASP scenario prompt:\n" + json.dumps(
+            prompt,
+            indent=2,
+            sort_keys=True,
         )
 
     def _finalize_terminal_episode(self, observation_record: dict[str, Any]) -> None:

server/scenario_cache.py

@@ -265,6 +265,29 @@ class ScenarioCache:
             counts[split][difficulty] = counts[split].get(difficulty, 0) + 1
         return {"root": str(self.root), "counts": counts, "entries": len(self._manifest_entries())}
 
+    def validated_entries(
+        self,
+        *,
+        split: str | None = None,
+        difficulty: int | None = None,
+    ) -> list[dict[str, Any]]:
+        entries = [
+            dict(entry)
+            for entry in self._manifest_entries()
+            if entry.get("validated") is True
+            and (split is None or entry.get("split") == split)
+            and (difficulty is None or int(entry.get("difficulty", -1)) == int(difficulty))
+        ]
+        return sorted(
+            entries,
+            key=lambda item: (
+                str(item.get("split", "")),
+                int(item.get("difficulty", 0)),
+                int(item.get("seed", 0)),
+                str(item.get("scenario_hash", "")),
+            ),
+        )
+
     def assert_coverage(self, *, split: str, difficulty: int | None = None) -> dict[str, Any]:
         coverage = self.coverage()
         required = self.settings.curriculum.minimum_for_split(split)
@@ -490,6 +513,8 @@ def _manifest_entry(
         "seed": int(scenario_record.get("seed", 0)),
         "split": str(scenario_record.get("split", "train")),
         "difficulty": int(scenario_record.get("difficulty", 0)),
+        "template_id": str(scenario_record.get("template_id", "")),
+        "bug_family": str(scenario_record.get("bug_family", "")),
         "scenario_hash": str(metadata.get("scenario_hash", "")),
         "cache_key": metadata.get("cache_key", {}),
         "validated": bool(metadata.get("validated", False)),

tests/test_closed_loop_runtime.py

@@ -49,6 +49,34 @@ def test_reset_records_scenario_family_and_partial_observability():
     assert "injected bug" not in serialized_hint
 
 
+def test_reset_returns_visible_scenario_prompt_without_hidden_identifiers():
+    env = make_env(75)
+    obs = env.reset(seed=75, split="train", difficulty=0)
+    prompt = obs.scenario_prompt
+    hidden = dict(env.state.hidden_facts)
+
+    assert "CyberSecurity_OWASP scenario prompt" in prompt
+    assert "available_actions" in prompt
+    assert str(env.state.seed) in prompt
+    assert env.state.scenario_hash in prompt
+    assert env.state.template_id in prompt
+
+    for key in (
+        "owner_user_id",
+        "intruder_user_id",
+        "admin_user_id",
+        "owner_invoice_id",
+        "other_invoice_id",
+        "foreign_invoice_id",
+        "tenant_a",
+        "tenant_b",
+    ):
+        value = str(hidden.get(key, ""))
+        assert not value or value not in prompt
+    assert "hidden_tests" not in prompt.lower()
+    assert "oracle" not in prompt.lower()
+
+
 def test_authz_oracle_fails_vulnerable_app_and_passes_secure_patch():
     env = make_env(72)
     oracle = AuthzOracle()

tests/test_grpo_curriculum.py

@@ -0,0 +1,138 @@
+from training.grpo_curriculum import (
+    AdaptiveDifficultyCurriculum,
+    ScenarioGroupRegistry,
+    build_scenario_group_rows,
+)
+
+
+def _entries():
+    return [
+        {
+            "seed": 10,
+            "split": "train",
+            "difficulty": 0,
+            "template_id": "fastapi_basic",
+            "bug_family": "bola_idor",
+            "scenario_hash": "hash-a",
+            "validated": True,
+        },
+        {
+            "seed": 20,
+            "split": "train",
+            "difficulty": 1,
+            "template_id": "fastapi_basic",
+            "bug_family": "bfla",
+            "scenario_hash": "hash-b",
+            "validated": True,
+        },
+        {
+            "seed": 30,
+            "split": "train",
+            "difficulty": 1,
+            "template_id": "fastapi_basic",
+            "bug_family": "tenant_leak",
+            "scenario_hash": "hash-c",
+            "validated": True,
+        },
+    ]
+
+
+def test_scenario_group_reuses_assignment_for_all_generations():
+    registry = ScenarioGroupRegistry(
+        _entries(),
+        split="train",
+        initial_difficulty=0,
+        rng_seed=1,
+        max_level=1,
+    )
+
+    first = registry.assignment_for(scenario_group_id=101, difficulty_policy="adaptive")
+    second = registry.assignment_for(scenario_group_id=101, difficulty_policy="adaptive")
+
+    assert first == second
+
+
+def test_different_scenario_groups_use_different_cached_scenarios_when_available():
+    registry = ScenarioGroupRegistry(
+        _entries(),
+        split="train",
+        initial_difficulty=1,
+        rng_seed=3,
+        max_level=1,
+    )
+
+    first = registry.assignment_for(
+        scenario_group_id=201,
+        requested_seed=20,
+        requested_difficulty=1,
+        split="train",
+        difficulty_policy="fixed",
+    )
+    second = registry.assignment_for(
+        scenario_group_id=202,
+        requested_seed=30,
+        requested_difficulty=1,
+        split="train",
+        difficulty_policy="fixed",
+    )
+
+    assert first["scenario_hash"] != second["scenario_hash"]
+
+
+def test_fixed_assignment_uses_dataset_seed_and_difficulty():
+    registry = ScenarioGroupRegistry(
+        _entries(),
+        split="train",
+        initial_difficulty=0,
+        rng_seed=1,
+        max_level=1,
+    )
+
+    assignment = registry.assignment_for(
+        scenario_group_id=301,
+        requested_seed=20,
+        requested_difficulty=1,
+        split="train",
+        difficulty_policy="fixed",
+    )
+
+    assert assignment["seed"] == 20
+    assert assignment["difficulty"] == 1
+    assert assignment["scenario_hash"] == "hash-b"
+
+
+def test_adaptive_curriculum_promotes_and_demotes_at_thresholds():
+    promote = AdaptiveDifficultyCurriculum(
+        min_level=0,
+        max_level=2,
+        current_level=0,
+        promote_after=50,
+    )
+    for _ in range(50):
+        promote.update(0, True)
+    assert promote.current_level == 1
+
+    demote = AdaptiveDifficultyCurriculum(
+        min_level=0,
+        max_level=2,
+        current_level=1,
+        promote_after=50,
+    )
+    for _ in range(50):
+        demote.update(1, False)
+    assert demote.current_level == 0
+
+
+def test_build_scenario_group_rows_include_grpo_group_columns():
+    rows = build_scenario_group_rows(
+        dataset_size=2,
+        training_prompt="repair local app",
+        seed_start=7,
+        split="train",
+        difficulty=1,
+    )
+
+    assert rows[0]["scenario_group_id"] == 7
+    assert rows[1]["scenario_group_id"] == 8
+    assert rows[0]["difficulty_policy"] == "adaptive"
+    assert rows[0]["prompt"][0]["content"] == "repair local app"

tests/test_trackio_utils.py

@@ -102,6 +102,8 @@ def test_trace_fingerprint_ignores_episode_id_but_tracks_action_changes():
         "scenario/split": "train",
         "scenario/difficulty": 0,
         "scenario/bug_type": "bola_idor",
+        "scenario/template_id": "fastapi_basic",
+        "scenario_hash": "scenario-a",
         "action_history": [
             {
                 "tool_name": "read_file",
@@ -113,11 +115,17 @@ def test_trace_fingerprint_ignores_episode_id_but_tracks_action_changes():
     }
     same_trace = dict(base_record)
     same_trace["episode_id"] = "episode-b"
+    token_only_reward_change = dict(base_record)
+    token_only_reward_change["reward_total"] = -0.25
     changed_trace = dict(base_record)
     changed_trace["action_history"] = [
         *base_record["action_history"],
         {"tool_name": "submit_fix", "arguments": {}},
     ]
+    different_scenario = dict(base_record)
+    different_scenario["scenario_hash"] = "scenario-b"
 
     assert episode_trace_fingerprint(base_record) == episode_trace_fingerprint(same_trace)
+    assert episode_trace_fingerprint(base_record) == episode_trace_fingerprint(token_only_reward_change)
     assert episode_trace_fingerprint(base_record) != episode_trace_fingerprint(changed_trace)
+    assert episode_trace_fingerprint(base_record) != episode_trace_fingerprint(different_scenario)

training/grpo_curriculum.py

@@ -0,0 +1,260 @@
+"""Scenario grouping and adaptive curriculum helpers for GRPO training."""
+
+from __future__ import annotations
+
+import random
+import threading
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class AdaptiveDifficultyCurriculum:
+    min_level: int = 0
+    max_level: int = 3
+    current_level: int = 0
+    promote_after: int = 50
+    promote_threshold: float = 0.70
+    demote_threshold: float = 0.35
+    ema_alpha: float = 0.10
+    rng_seed: int = 0
+    counts: dict[int, int] = field(default_factory=dict)
+    ema_success: dict[int, float] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        self.min_level = int(self.min_level)
+        self.max_level = int(self.max_level)
+        self.current_level = max(self.min_level, min(int(self.current_level), self.max_level))
+        self._rng = random.Random(int(self.rng_seed))
+
+    def sample_difficulty(self, available_difficulties: Iterable[int]) -> int:
+        available = {int(item) for item in available_difficulties}
+        if not available:
+            raise ValueError("No cached difficulties are available for GRPO curriculum sampling.")
+
+        candidates = [
+            max(self.min_level, self.current_level - 1),
+            self.current_level,
+            min(self.max_level, self.current_level + 1),
+        ]
+        weights = [0.20, 0.65, 0.15]
+        weighted: dict[int, float] = {}
+        for level, weight in zip(candidates, weights):
+            if level in available:
+                weighted[level] = weighted.get(level, 0.0) + weight
+
+        if not weighted:
+            nearest = min(available, key=lambda level: (abs(level - self.current_level), level))
+            return nearest
+        levels = list(weighted)
+        return int(self._rng.choices(levels, weights=[weighted[level] for level in levels], k=1)[0])
+
+    def update(self, difficulty: int, success: float | bool) -> dict[str, Any]:
+        level = int(difficulty)
+        value = max(0.0, min(1.0, float(success)))
+        self.counts[level] = self.counts.get(level, 0) + 1
+        old = self.ema_success.get(level, 0.0)
+        self.ema_success[level] = (1.0 - self.ema_alpha) * old + self.ema_alpha * value
+
+        if level == self.current_level and self.counts[level] >= self.promote_after:
+            if self.ema_success[level] >= self.promote_threshold:
+                self.current_level = min(self.max_level, self.current_level + 1)
+            elif self.ema_success[level] <= self.demote_threshold:
+                self.current_level = max(self.min_level, self.current_level - 1)
+        return self.snapshot()
+
+    def snapshot(self) -> dict[str, Any]:
+        return {
+            "current_level": self.current_level,
+            "counts": {str(key): value for key, value in sorted(self.counts.items())},
+            "ema_success": {
+                str(key): value for key, value in sorted(self.ema_success.items())
+            },
+            "current_level_ema_success": self.ema_success.get(self.current_level, 0.0),
+        }
+
+
+class ScenarioGroupRegistry:
+    """Assign each GRPO group to exactly one cached scenario."""
+
+    def __init__(
+        self,
+        entries: Sequence[Mapping[str, Any]],
+        *,
+        split: str = "train",
+        initial_difficulty: int = 0,
+        rng_seed: int = 0,
+        max_level: int | None = None,
+    ) -> None:
+        self.split = split
+        self._rng = random.Random(int(rng_seed))
+        self._lock = threading.Lock()
+        self._assignments: dict[int, dict[str, Any]] = {}
+        self._completed_groups: set[int] = set()
+        self._entries_by_difficulty: dict[int, list[dict[str, Any]]] = {}
+        self._cursors: dict[int, int] = {}
+
+        for entry in entries:
+            if entry.get("validated") is not True or entry.get("split") != split:
+                continue
+            difficulty = int(entry.get("difficulty", 0))
+            self._entries_by_difficulty.setdefault(difficulty, []).append(dict(entry))
+
+        for difficulty, items in self._entries_by_difficulty.items():
+            items.sort(key=lambda item: (int(item.get("seed", 0)), str(item.get("scenario_hash", ""))))
+            self._rng.shuffle(items)
+            self._cursors[difficulty] = 0
+
+        if not self._entries_by_difficulty:
+            raise ValueError(f"No validated cached scenarios are available for split={split!r}.")
+
+        available = sorted(self._entries_by_difficulty)
+        resolved_max = max_level if max_level is not None else max(available)
+        self.curriculum = AdaptiveDifficultyCurriculum(
+            min_level=min(available),
+            max_level=int(resolved_max),
+            current_level=int(initial_difficulty),
+            rng_seed=int(rng_seed),
+        )
+
+    @property
+    def available_difficulties(self) -> list[int]:
+        return sorted(self._entries_by_difficulty)
+
+    def assignment_for(
+        self,
+        *,
+        scenario_group_id: int,
+        requested_seed: int | None = None,
+        requested_difficulty: int | None = None,
+        split: str | None = None,
+        difficulty_policy: str = "adaptive",
+    ) -> dict[str, Any]:
+        group_id = int(scenario_group_id)
+        with self._lock:
+            if group_id in self._assignments:
+                return dict(self._assignments[group_id])
+
+            if difficulty_policy == "fixed":
+                difficulty = int(
+                    requested_difficulty
+                    if requested_difficulty is not None
+                    else self.curriculum.current_level
+                )
+                entry = self._find_entry(
+                    seed=requested_seed,
+                    split=split or self.split,
+                    difficulty=difficulty,
+                ) or self._next_entry(difficulty)
+            else:
+                difficulty = self.curriculum.sample_difficulty(self.available_difficulties)
+                entry = self._next_entry(difficulty)
+
+            assignment = self._assignment_from_entry(group_id, entry)
+            self._assignments[group_id] = assignment
+            return dict(assignment)
+
+    def record_group_outcome(self, scenario_group_id: int, success_rate: float) -> dict[str, Any] | None:
+        group_id = int(scenario_group_id)
+        with self._lock:
+            if group_id in self._completed_groups:
+                return None
+            self._completed_groups.add(group_id)
+            assignment = self._assignments.get(group_id)
+            if not assignment:
+                return self.curriculum.snapshot()
+            return self.curriculum.update(
+                int(assignment["difficulty"]),
+                max(0.0, min(1.0, float(success_rate))),
+            )
+
+    def metrics(
+        self,
+        records: Sequence[Mapping[str, Any]],
+        *,
+        unique_trace_count: int,
+        duplicate_trace_suppressed_count: int,
+    ) -> dict[str, float]:
+        scenario_hashes = {
+            str(record.get("scenario_hash") or record.get("scenario_id_hash") or "")
+            for record in records
+            if record.get("scenario_hash") or record.get("scenario_id_hash")
+        }
+        seeds = {
+            int(record.get("scenario/seed", record.get("seed", 0)) or 0)
+            for record in records
+        }
+        total = max(1, len(records))
+        snapshot = self.curriculum.snapshot()
+        return {
+            "train/unique_trace_count": float(unique_trace_count),
+            "train/duplicate_trace_suppressed_count": float(duplicate_trace_suppressed_count),
+            "train/unique_trace_rate": float(unique_trace_count) / total,
+            "train/unique_seed_count": float(len(seeds)),
+            "train/unique_scenario_hash_count": float(len(scenario_hashes)),
+            "train/curriculum_level": float(snapshot["current_level"]),
+            "train/curriculum_ema_success": float(snapshot["current_level_ema_success"]),
+        }
+
+    def _find_entry(
+        self,
+        *,
+        seed: int | None,
+        split: str,
+        difficulty: int,
+    ) -> dict[str, Any] | None:
+        if seed is None or split != self.split:
+            return None
+        for entry in self._entries_by_difficulty.get(int(difficulty), []):
+            if int(entry.get("seed", -1)) == int(seed):
+                return dict(entry)
+        return None
+
+    def _next_entry(self, difficulty: int) -> dict[str, Any]:
+        level = int(difficulty)
+        items = self._entries_by_difficulty.get(level)
+        if not items:
+            nearest = min(
+                self.available_difficulties,
+                key=lambda item: (abs(item - level), item),
+            )
+            items = self._entries_by_difficulty[nearest]
+            level = nearest
+        cursor = self._cursors.get(level, 0)
+        self._cursors[level] = cursor + 1
+        return dict(items[cursor % len(items)])
+
+    def _assignment_from_entry(self, group_id: int, entry: Mapping[str, Any]) -> dict[str, Any]:
+        cache_key = entry.get("cache_key") if isinstance(entry.get("cache_key"), Mapping) else {}
+        return {
+            "scenario_group_id": int(group_id),
+            "seed": int(entry.get("seed", 0)),
+            "split": str(entry.get("split", self.split)),
+            "difficulty": int(entry.get("difficulty", 0)),
+            "scenario_hash": str(entry.get("scenario_hash", "")),
+            "template_id": str(entry.get("template_id") or cache_key.get("app_family", "")),
+            "bug_family": str(entry.get("bug_family") or cache_key.get("authz_bug_type", "")),
+        }
+
+
+def build_scenario_group_rows(
+    *,
+    dataset_size: int,
+    training_prompt: str,
+    seed_start: int = 0,
+    split: str = "train",
+    difficulty: int = 0,
+    difficulty_policy: str = "adaptive",
+) -> list[dict[str, Any]]:
+    return [
+        {
+            "prompt": [{"role": "user", "content": training_prompt}],
+            "scenario_group_id": int(seed_start) + index,
+            "seed": int(seed_start) + index,
+            "difficulty": int(difficulty),
+            "split": split,
+            "difficulty_policy": difficulty_policy,
+        }
+        for index in range(int(dataset_size))
+    ]

training/trackio_utils.py

@@ -150,9 +150,16 @@ REQUIRED_SMOKE_TRACKIO_ITEMS = (
 TRACE_TABLE_COLUMNS = (
     "episode_id",
     "scenario_id_hash",
+    "scenario_hash",
+    "seed",
     "split",
     "difficulty",
+    "template_id",
     "bug_type",
+    "reward_total",
+    "security_pass_rate",
+    "regression_pass_rate",
+    "step_count",
     "visible_observation_summary",
     "action_sequence",
     "tool_calls",
@@ -529,6 +536,7 @@ def episode_record_from_state(
         "target_weakness": getattr(state, "target_weakness", ""),
         "difficulty_tier": getattr(state, "difficulty_tier", ""),
         "domain": getattr(state, "domain", ""),
+        "scenario_hash": getattr(state, "scenario_hash", ""),
         "success": bool(getattr(state, "success", False)),
         "failure_reason": getattr(state, "failure_reason", None),
         "finding_submitted": bool(getattr(state, "finding_submitted", False)),
@@ -821,12 +829,20 @@ def episode_to_trace_row(episode: Any) -> dict[str, Any]:
     files_modified = _files_modified(record, actions)
     reward_breakdown = _final_reward_breakdown(record)
     final_obs = observations[-1] if observations else {}
+    tracking_fields = episode_to_tracking_fields(record)
     row = {
         "episode_id": _redact_text(record.get("episode_id", "")),
         "scenario_id_hash": record.get("scenario_id_hash") or _scenario_hash(record),
+        "scenario_hash": record.get("scenario_hash") or _as_dict(record.get("metrics")).get("scenario_hash", ""),
+        "seed": record.get("scenario/seed") or record.get("seed", 0),
         "split": record.get("scenario/split") or record.get("split", ""),
         "difficulty": record.get("scenario/difficulty") or record.get("difficulty", 0),
+        "template_id": record.get("scenario/template_id") or record.get("template_id", ""),
         "bug_type": record.get("scenario/bug_type") or record.get("bug_type", ""),
+        "reward_total": tracking_fields["reward/total"],
+        "security_pass_rate": tracking_fields["reward/hidden_authz_pass_rate"],
+        "regression_pass_rate": tracking_fields["reward/normal_flow_pass_rate"],
+        "step_count": record.get("step_count", len(actions)),
         "visible_observation_summary": json.dumps(
             {
                 "done": bool(record.get("done", final_obs.get("done", False))),
@@ -845,9 +861,7 @@ def episode_to_trace_row(episode: Any) -> dict[str, Any]:
                 "local_probe_count": sum(
                     1 for name in tool_names if name in {"send_local_request", "compare_identities"}
                 ),
-                "first_valid_exploit_step": episode_to_tracking_fields(record)[
-                    "skill/first_valid_exploit_step"
-                ],
+                "first_valid_exploit_step": tracking_fields["skill/first_valid_exploit_step"],
                 "diagnosis_submitted": bool(
                     record.get("diagnosis_submitted", record.get("finding_submitted", False))
                 ),
@@ -894,7 +908,7 @@ def episode_trace_fingerprint(episode: Any) -> str:
         {
             key: row.get(key, "")
             for key in TRACE_TABLE_COLUMNS
-            if key != "episode_id"
+            if key not in {"episode_id", "reward_total"}
         },
         length=24,
     )