Sync current main to Hugging Face Space

README.md

@@ -174,9 +174,27 @@ osint-env train-self-play --config config/shared_config.json --train-config conf
 
 When you have compute and the train dependencies installed, remove `--dry-run` (or set `"dry_run": false` in the training config) to execute TRL GRPO updates for alternating generator and answerer phases.
 
+For a standalone Linux server or SSH box, there is also a wrapper script that activates a venv, optionally installs train deps, and runs the same command:
+
+```bash
+VENV_PATH="$HOME/arl" \
+INSTALL_TRAIN_DEPS=1 \
+TRAIN_ENV_CONFIG_PATH="config/shared_config.json" \
+TRAIN_SELF_PLAY_CONFIG_PATH="config/self_play_training_hf_a10g_smoke.json" \
+TRAIN_SELF_PLAY_OUTPUT_DIR="artifacts/self_play_server" \
+bash scripts/train_self_play_standalone.sh
+```
+
+Useful overrides for the standalone script:
+
+- `BOOTSTRAP_VENV=1` to create the virtualenv if it does not exist
+- `TRAIN_SELF_PLAY_ROUNDS=2` to override the number of rounds
+- `RUN_SELF_PLAY_DRY_RUN=1` to materialize artifacts without GRPO updates
+- `TRAIN_SETUP_COMMAND='python -m pip install flash-attn --no-build-isolation'` for host-specific extras
+
 The training config also supports `"model_topology": "dual"|"shared"`, `"phase_schedule": "generator_answerer"|"answerer_generator_answerer"`, `"tuning_mode": "full"|"lora"`, and `"canonical_graph_mode": "generate"|"fixed"` so you can switch between two-model vs single-model self-play, full fine-tuning vs LoRA adapters, and whether canonical graph structure is generated each round or kept fixed while training question/answer behavior.
 
-### Hugging Face Space Smoke Run (Qwen 3.5 0.8B + W&B)
+### Hugging Face Job A10G Run (Separate From The Space)
 
 For a short verification run (enough to confirm W&B logging before scaling up), use:
 
@@ -186,27 +204,42 @@ osint-env train-self-play --config config/shared_config.json --train-config conf
 
 This config:
 
-- uses `Qwen/Qwen3.5-0.8B`
+- uses `Qwen/Qwen2.5-0.5B-Instruct`
 - enables W&B reporting (`wandb_enabled: true`)
 - uses `pipeline_mode: "swarm_v2"` with `canonical_graph_mode: "fixed"` to keep canonical graph candidates stable while training question/answer behavior
-- keeps training intentionally short (`rounds=1`, `max_steps=5` per phase)
-- uses LoRA with small batch settings so it can run as a smoke test on an A10G
+- keeps training intentionally short (`rounds=2`, `max_steps=50` per phase)
+- uses full fine-tuning plus fused AdamW, bf16/tf32, larger generation batches, and extra dataloader workers to make better use of an A10G
 
 To enable canonical graph generation during swarm_v2 training, switch `"canonical_graph_mode"` to `"generate"` in the training config.
 
-Space setup checklist:
+If you want the Space to stay on CPU and train separately on paid GPU compute, launch a dedicated Hugging Face Job instead of training inside the Space:
+
+```bash
+osint-env-launch-hf-job \
+  --hf-token "$HF_TOKEN" \
+  --job-image "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel" \
+  --repo-url "https://github.com/your-org/meta-knowledge-graph.git" \
+  --repo-ref "main" \
+  --flavor "a10g-small" \
+  --env-config "config/shared_config.json" \
+  --train-config "config/self_play_training_hf_a10g_smoke.json" \
+  --output-bucket "your-hf-bucket" \
+  --wait
+```
+
+The launcher talks to the Hugging Face Jobs API through `huggingface_hub`, so the Space can remain on CPU while the training job runs on separate A10G compute.
+
+Optional Space startup wiring still exists if you want it:
 
-1. In Space **Settings -> Hardware**, select **NVIDIA A10G (large)**.
-2. In Space **Settings -> Variables and secrets**, set `WANDB_API_KEY`.
-3. Set `HF_TOKEN` in Space secrets to avoid unauthenticated Hub downloads and stricter rate limits.
-4. Optionally set `WANDB_ENTITY` if your project belongs to a team.
-5. Set `RUN_SELF_PLAY_TRAINING=1` in Space variables to trigger training during container startup.
-6. Optional overrides:
+1. Keep the Space on CPU if it is serving inference/UI only.
+2. Set `RUN_SELF_PLAY_TRAINING=1` only if you intentionally want startup-time training inside the Space container.
+3. Optional overrides:
    - `TRAIN_SELF_PLAY_CONFIG_PATH` (default: `config/self_play_training_hf_a10g_smoke.json`)
    - `TRAIN_ENV_CONFIG_PATH` (default: `config/shared_config.json`)
-   - `RUN_SELF_PLAY_DRY_RUN=1` to test startup wiring without GRPO updates.
-   - `OSINT_TRAIN_STRICT_ASSERTS=1` to fail fast when reward variance, KL, loss, grad norms, or parameter updates stay zero.
-7. Restart the Space and monitor build/runtime logs for the training run.
+   - `TRAIN_SELF_PLAY_OUTPUT_DIR` to override where artifacts land
+   - `RUN_SELF_PLAY_DRY_RUN=1` to test startup wiring without GRPO updates
+   - `RUN_SELF_PLAY_BACKGROUND=1` to keep the API up while startup-time training runs
+   - `OSINT_TRAIN_STRICT_ASSERTS=1` to fail fast when reward variance, KL, loss, grad norms, or parameter updates stay zero
 
 W&B run naming is controlled by `wandb_run_name_prefix` and will emit phase-specific runs like `...-r001-generator` and `...-r001-answerer`.
 
@@ -229,10 +262,10 @@ In `legacy` pipeline mode, the reward is a weighted sum:
 
 Default weights (configurable through `generator_reward_weights` in training config):
 
-- `validity`: `0.35`
-- `hardness`: `0.45`
-- `diversity`: `0.10`
-- `consistency`: `0.10`
+- `validity`: `0.45`
+- `hardness`: `0.20`
+- `diversity`: `0.15`
+- `consistency`: `0.20`
 
 In `swarm_v2` pipeline mode, generation uses strict replay/validation first, then a structured reward:
 

config/self_play_training_example.json

@@ -15,11 +15,14 @@
   "max_graph_context_edges": 100,
   "max_support_edges": 8,
   "answerer_judge_max_new_tokens": 48,
+  "generated_task_max_new_tokens": 512,
+  "post_training_eval_questions": 24,
+  "post_training_eval_answer_max_new_tokens": 128,
   "generator_reward_weights": {
-    "validity": 0.35,
-    "hardness": 0.45,
-    "diversity": 0.1,
-    "consistency": 0.1
+    "validity": 0.45,
+    "hardness": 0.2,
+    "diversity": 0.15,
+    "consistency": 0.2
   },
   "lora": {
     "r": 16,
@@ -62,12 +65,14 @@
   },
   "generator_phase": {
     "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
-    "learning_rate": 1e-06,
+    "learning_rate": 5e-06,
     "max_steps": 64,
     "per_device_train_batch_size": 2,
     "gradient_accumulation_steps": 4,
     "num_generations": 4,
-    "max_completion_length": 256,
+    "max_completion_length": 384,
+    "max_prompt_length": 1024,
+    "generation_batch_size": 8,
     "temperature": 1.0,
     "top_p": 1.0,
     "beta": 0.01,
@@ -77,18 +82,28 @@
     "scale_rewards": "none",
     "logging_steps": 10,
     "save_steps": 50,
+    "save_total_limit": 2,
+    "optim": "adamw_torch_fused",
+    "bf16": true,
+    "tf32": true,
+    "gradient_checkpointing": false,
+    "dataloader_num_workers": 2,
+    "dataloader_persistent_workers": true,
+    "dataloader_prefetch_factor": 2,
     "output_subdir": "generator",
     "use_vllm": false,
     "vllm_mode": "colocate"
   },
   "answerer_phase": {
     "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
-    "learning_rate": 1e-06,
+    "learning_rate": 3e-06,
     "max_steps": 64,
     "per_device_train_batch_size": 2,
     "gradient_accumulation_steps": 4,
     "num_generations": 4,
     "max_completion_length": 192,
+    "max_prompt_length": 1024,
+    "generation_batch_size": 8,
     "temperature": 1.0,
     "top_p": 1.0,
     "beta": 0.01,
@@ -98,6 +113,14 @@
     "scale_rewards": "none",
     "logging_steps": 10,
     "save_steps": 50,
+    "save_total_limit": 2,
+    "optim": "adamw_torch_fused",
+    "bf16": true,
+    "tf32": true,
+    "gradient_checkpointing": false,
+    "dataloader_num_workers": 2,
+    "dataloader_persistent_workers": true,
+    "dataloader_prefetch_factor": 2,
     "output_subdir": "answerer",
     "use_vllm": false,
     "vllm_mode": "colocate"

config/self_play_training_hf_a10g_smoke.json

@@ -10,7 +10,7 @@
   "canonical_graph_mode": "fixed",
   "model_topology": "shared",
   "phase_schedule": "generator_answerer",
-  "tuning_mode": "lora",
+  "tuning_mode": "full",
   "shared_model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
   "seed_tasks_per_round": 16,
   "generated_tasks_per_round": 24,
@@ -19,14 +19,19 @@
   "max_graph_context_edges": 24,
   "max_support_edges": 6,
   "answerer_judge_max_new_tokens": 32,
+  "generated_task_max_new_tokens": 640,
+  "post_training_eval_questions": 24,
+  "post_training_eval_answer_max_new_tokens": 128,
   "generator_phase": {
     "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
-    "learning_rate": 1e-06,
+    "learning_rate": 5e-06,
     "max_steps": 50,
     "per_device_train_batch_size": 4,
     "gradient_accumulation_steps": 1,
     "num_generations": 4,
-    "max_completion_length": 768,
+    "max_completion_length": 384,
+    "max_prompt_length": 768,
+    "generation_batch_size": 16,
     "temperature": 0.9,
     "top_p": 0.95,
     "repetition_penalty": 1.1,
@@ -37,18 +42,28 @@
     "scale_rewards": "group",
     "logging_steps": 1,
     "save_steps": 10,
+    "save_total_limit": 2,
+    "optim": "adamw_torch_fused",
+    "bf16": true,
+    "tf32": true,
+    "gradient_checkpointing": false,
+    "dataloader_num_workers": 4,
+    "dataloader_persistent_workers": true,
+    "dataloader_prefetch_factor": 4,
     "output_subdir": "generator_train",
     "use_vllm": false,
     "vllm_mode": "colocate"
   },
   "answerer_phase": {
     "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
-    "learning_rate": 1e-06,
+    "learning_rate": 3e-06,
     "max_steps": 50,
     "per_device_train_batch_size": 4,
     "gradient_accumulation_steps": 1,
     "num_generations": 4,
     "max_completion_length": 256,
+    "max_prompt_length": 768,
+    "generation_batch_size": 16,
     "temperature": 0.7,
     "top_p": 0.95,
     "repetition_penalty": 1.1,
@@ -59,6 +74,14 @@
     "scale_rewards": "group",
     "logging_steps": 1,
     "save_steps": 10,
+    "save_total_limit": 2,
+    "optim": "adamw_torch_fused",
+    "bf16": true,
+    "tf32": true,
+    "gradient_checkpointing": false,
+    "dataloader_num_workers": 4,
+    "dataloader_persistent_workers": true,
+    "dataloader_prefetch_factor": 4,
     "output_subdir": "answerer_train",
     "use_vllm": false,
     "vllm_mode": "colocate"

docs/adversarial_self_play.md

@@ -59,6 +59,14 @@ This directly supports the "train solver, freeze, attack, retrain solver" sequen
 
 Weights are configurable in `generator_reward_weights`.
 
+For `swarm_v2`, the reward now prioritizes:
+
+- Valid, replayable task structure first.
+- Hardness against the frozen answerer second.
+- Diversity and compact multi-agent/shared-context usage after validity.
+
+This avoids the degenerate regime where almost every sample is invalid and the whole batch stays negative.
+
 ### Answerer (existing reward integration)
 
 `AnswererRewardFunction` wraps existing environment reward logic:
@@ -88,12 +96,60 @@ In dry run mode, the pipeline still:
 
 But it skips expensive GRPO updates.
 
+## Post-Training Evaluation
+
+After a non-dry-run training job completes, the runner now writes a post-training evaluation artifact that:
+
+- Uses the finetuned generator to create fresh evaluation questions.
+- Evaluates both the finetuned answerer and the original/base answerer on those generated questions.
+- Reports a `delta_vs_original` summary so you can see whether fine-tuning actually improved task success, reward, and graph F1.
+- Saves the summary and episode rows under `post_training_evaluation.json`.
+
+You can control this flow with:
+
+- `generated_task_max_new_tokens`: decoding budget for generator-side task sampling/eval.
+- `post_training_eval_questions`: how many fresh tasks to evaluate after training.
+- `post_training_eval_answer_max_new_tokens`: answerer decoding budget for the final eval pass.
+
+## Checkpoints And Final Models
+
+Self-play outputs are written under `output_dir` (default `artifacts/self_play`) unless overridden.
+
+Per round and phase you will now find:
+
+- `round_XXX/<phase>/checkpoint-*`: intermediate trainer checkpoints saved every `save_steps`.
+- `round_XXX/<phase>/final_model`: final saved model for that phase, with tokenizer files.
+- `self_play_summary.json`: top-level run summary.
+- `post_training_evaluation.json`: generated-question evaluation written after training.
+
 ## Compute Mode
 
 When compute is available:
 
 1. Install train dependencies: `python -m pip install -e ".[train]"`
 2. Disable dry run (`--dry-run` off and/or `"dry_run": false` in config).
-3. Run `osint-env train-self-play`.
+3. Run `osint-env train-self-play`, or launch a dedicated Hugging Face Job with `osint-env-launch-hf-job` if you want the Space to stay on CPU while training runs on separate GPU compute.
 
 Outputs are written under `artifacts/self_play` unless overridden.
+
+## Standalone Server Script
+
+For an SSH server or other standalone machine, you can use `scripts/train_self_play_standalone.sh`.
+
+Example:
+
+```bash
+VENV_PATH="$HOME/arl" \
+INSTALL_TRAIN_DEPS=1 \
+TRAIN_ENV_CONFIG_PATH="config/shared_config.json" \
+TRAIN_SELF_PLAY_CONFIG_PATH="config/self_play_training_hf_a10g_smoke.json" \
+TRAIN_SELF_PLAY_OUTPUT_DIR="artifacts/self_play_server" \
+bash scripts/train_self_play_standalone.sh
+```
+
+Useful environment variables:
+
+- `BOOTSTRAP_VENV=1`: create the virtualenv automatically if it does not exist yet.
+- `TRAIN_SELF_PLAY_ROUNDS=2`: override the round count without editing JSON.
+- `RUN_SELF_PLAY_DRY_RUN=1`: skip GRPO updates and only materialize artifacts.
+- `TRAIN_SETUP_COMMAND='python -m pip install flash-attn --no-build-isolation'`: run any host-specific setup before training.

pyproject.toml

@@ -22,6 +22,7 @@ train = [
 	"accelerate>=0.33.0",
 	"trl>=0.15.0",
 	"peft>=0.11.0",
+	"huggingface_hub>=0.34.0",
 	"pillow",
 	"torchvision",
 	"wandb",
@@ -30,6 +31,7 @@ train = [
 [project.scripts]
 osint-env = "osint_env.cli:main"
 server = "osint_env.server_entry:main"
+osint-env-launch-hf-job = "osint_env.training.hf_jobs:main"
 
 [build-system]
 requires = ["setuptools>=68", "wheel"]

scripts/space_start.sh

@@ -10,23 +10,46 @@ _is_true() {
 
 ENV_CONFIG_PATH="${TRAIN_ENV_CONFIG_PATH:-config/shared_config.json}"
 TRAIN_CONFIG_PATH="${TRAIN_SELF_PLAY_CONFIG_PATH:-config/self_play_training_hf_a10g_smoke.json}"
+TRAIN_OUTPUT_DIR="${TRAIN_SELF_PLAY_OUTPUT_DIR:-}"
 RUN_FLAG="${RUN_SELF_PLAY_TRAINING:-0}"
 DRY_RUN_FLAG="${RUN_SELF_PLAY_DRY_RUN:-0}"
+BACKGROUND_FLAG="${RUN_SELF_PLAY_BACKGROUND:-1}"
+
+_train_self_play() {
+  if [ -n "${TRAIN_OUTPUT_DIR}" ]; then
+    OUTPUT_ARG="--train-output-dir ${TRAIN_OUTPUT_DIR}"
+  else
+    OUTPUT_ARG=""
+  fi
 
-if _is_true "$RUN_FLAG"; then
-  echo "[space_start] RUN_SELF_PLAY_TRAINING enabled."
-  echo "[space_start] Training start: $(date -u +"%Y-%m-%dT%H:%M:%SZ")"
-  echo "[space_start] Env config: ${ENV_CONFIG_PATH}"
-  echo "[space_start] Train config: ${TRAIN_CONFIG_PATH}"
   if _is_true "$DRY_RUN_FLAG"; then
     echo "[space_start] Running self-play in dry-run mode."
-    osint-env train-self-play --config "${ENV_CONFIG_PATH}" --train-config "${TRAIN_CONFIG_PATH}" --dry-run
+    # shellcheck disable=SC2086
+    osint-env train-self-play --config "${ENV_CONFIG_PATH}" --train-config "${TRAIN_CONFIG_PATH}" ${OUTPUT_ARG} --dry-run
   else
     echo "[space_start] Running self-play training."
-    osint-env train-self-play --config "${ENV_CONFIG_PATH}" --train-config "${TRAIN_CONFIG_PATH}"
+    # shellcheck disable=SC2086
+    osint-env train-self-play --config "${ENV_CONFIG_PATH}" --train-config "${TRAIN_CONFIG_PATH}" ${OUTPUT_ARG}
   fi
+
   echo "[space_start] Self-play command completed."
   echo "[space_start] Training end: $(date -u +"%Y-%m-%dT%H:%M:%SZ")"
+}
+
+if _is_true "$RUN_FLAG"; then
+  echo "[space_start] RUN_SELF_PLAY_TRAINING enabled."
+  echo "[space_start] Training start: $(date -u +"%Y-%m-%dT%H:%M:%SZ")"
+  echo "[space_start] Env config: ${ENV_CONFIG_PATH}"
+  echo "[space_start] Train config: ${TRAIN_CONFIG_PATH}"
+  if [ -n "${TRAIN_OUTPUT_DIR}" ]; then
+    echo "[space_start] Train output dir: ${TRAIN_OUTPUT_DIR}"
+  fi
+  if _is_true "$BACKGROUND_FLAG"; then
+    echo "[space_start] Launching self-play in background so the Space API can stay online."
+    _train_self_play &
+  else
+    _train_self_play
+  fi
 else
   echo "[space_start] RUN_SELF_PLAY_TRAINING disabled. Skipping self-play run."
 fi

src/osint_env/agents/single_agent.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import re
+
 from osint_env.domain.models import Action, ActionType
 from osint_env.env.environment import OSINTEnvironment
 from osint_env.llm.interface import LLMClient, RuleBasedMockLLM
+from osint_env.platforms.tool_schemas import build_lookup_tools
 
 
 class SingleAgentRunner:
@@ -15,14 +18,31 @@ class SingleAgentRunner:
         done = False
         info = {}
         while not done:
-            messages = [{"role": "system", "content": f"question: {obs.task['question']}"}]
-            tools = []
+            messages = [
+                {
+                    "role": "system",
+                    "content": (
+                        f"question: {obs.task['question']}\n"
+                        f"shared_context_available: {bool(obs.task.get('shared_context_available', False))}\n"
+                        "Use lookup tools to gather evidence before answering."
+                    ),
+                }
+            ]
+            tools = build_lookup_tools()
             try:
                 llm_resp = self.llm.generate(messages, tools)
                 planned_calls = llm_resp.tool_calls[:2]
             except Exception:
                 planned_calls = []
 
+            if not planned_calls and bool(obs.task.get("shared_context_available", False)):
+                planned_calls = [
+                    {
+                        "tool_name": "search_shared_context",
+                        "args": {"query": self._shared_context_query(obs.task["question"]), "k": 5},
+                    }
+                ]
+
             for call in planned_calls:
                 obs, _, done, info = self.env.step(Action(ActionType.CALL_TOOL, call))
                 if done:
@@ -39,3 +59,16 @@ class SingleAgentRunner:
             if token.startswith("alias_") or token.startswith("user_"):
                 return token
         return "unknown"
+
+    @staticmethod
+    def _shared_context_query(question: str) -> str:
+        id_match = re.search(r"\b(?:alias|user|post|thr|thread|org|loc|event)_[A-Za-z0-9_]+\b", question)
+        if id_match:
+            return id_match.group(0)
+        path_match = re.search(r"relation path\s+(.+?),\s*which entity", question, flags=re.IGNORECASE)
+        if path_match:
+            first_relation = path_match.group(1).split("->", 1)[0].strip()
+            if first_relation:
+                return first_relation
+        tokens = re.findall(r"[A-Za-z0-9_]+", question)
+        return tokens[0] if tokens else question

src/osint_env/agents/swarm_agent.py

@@ -7,6 +7,7 @@ from osint_env.domain.models import Action, ActionType
 from osint_env.env.environment import OSINTEnvironment
 from osint_env.env.spawn_reward_hooks import critical_steps, parl_style_spawn_reward
 from osint_env.llm.interface import LLMClient, RuleBasedMockLLM
+from osint_env.platforms.tool_schemas import build_lookup_tools
 
 
 class SwarmAgentRunner:
@@ -135,12 +136,13 @@ class SwarmAgentRunner:
                 "content": (
                     f"question: {obs.task['question']}\n"
                     f"agent_role: {role}_{agent_idx}\n"
+                    f"shared_context_available: {bool(obs.task.get('shared_context_available', False))}\n"
                     "Return concise tool plan."
                 ),
             }
         ]
         try:
-            response = self.llm.generate(messages, tools=[])
+            response = self.llm.generate(messages, tools=build_lookup_tools())
         except Exception:
             response = None
 
@@ -160,6 +162,11 @@ class SwarmAgentRunner:
             return calls
 
         question = str(obs.task.get("question", "")).lower()
+        shared_context_available = bool(obs.task.get("shared_context_available", False))
+        shared_query = self._shared_context_query(str(obs.task.get("question", "")))
+        if shared_context_available and role in {"explorer", "reasoner"}:
+            return [{"tool_name": "search_shared_context", "args": {"query": shared_query, "k": 5}}]
+
         if role == "explorer":
             if "event" in question:
                 return [{"tool_name": "search_threads", "args": {"topic": "security"}}]
@@ -182,6 +189,19 @@ class SwarmAgentRunner:
 
         return [{"tool_name": "search_people", "args": {"org": "Apex"}}]
 
+    @staticmethod
+    def _shared_context_query(question: str) -> str:
+        id_match = re.search(r"\b(?:alias|user|post|thr|thread|org|loc|event)_[A-Za-z0-9_]+\b", question)
+        if id_match:
+            return id_match.group(0)
+        path_match = re.search(r"relation path\s+(.+?),\s*which entity", question, flags=re.IGNORECASE)
+        if path_match:
+            first_relation = path_match.group(1).split("->", 1)[0].strip()
+            if first_relation:
+                return first_relation
+        tokens = re.findall(r"[A-Za-z0-9_]+", question)
+        return tokens[0] if tokens else question
+
     def _edge_plan(self, agent_idx: int) -> dict[str, Any] | None:
         if self.env.state is None or not self.env.state.task.supporting_edges:
             return None

src/osint_env/baselines/openai_runner.py

@@ -12,6 +12,7 @@ from osint_env.env.environment import OSINTEnvironment
 from osint_env.env.reward import compute_graph_f1
 from osint_env.eval.leaderboard import append_leaderboard_record, load_leaderboard
 from osint_env.eval.metrics import EvalMetrics
+from osint_env.platforms.tool_schemas import build_action_tools
 from osint_env.viz import export_dashboard
 
 
@@ -50,123 +51,6 @@ class OpenAIBaselineConfig:
     max_steps: int = 8
     seed: int | None = 7
     append_leaderboard: bool = True
-
-
-def _tool_schema(
-    name: str,
-    description: str,
-    properties: dict[str, Any],
-    required: list[str],
-) -> dict[str, Any]:
-    return {
-        "type": "function",
-        "function": {
-            "name": name,
-            "description": description,
-            "parameters": {
-                "type": "object",
-                "properties": properties,
-                "required": required,
-                "additionalProperties": False,
-            },
-        },
-    }
-
-
-def build_action_tools() -> list[dict[str, Any]]:
-    return [
-        _tool_schema(
-            "search_posts",
-            "Search microblog posts by substring over post text, post id, author id, canonical user id, or referenced entity ids/names.",
-            {"query": {"type": "string", "description": "Substring to search for in post text."}},
-            ["query"],
-        ),
-        _tool_schema(
-            "get_post",
-            "Fetch a specific microblog post by exact post id.",
-            {"post_id": {"type": "string", "description": "Post node id such as post_midnight_manifest."}},
-            ["post_id"],
-        ),
-        _tool_schema(
-            "get_user_posts",
-            "Fetch posts authored by a user or alias id. Alias ids are resolved to the canonical user and vice versa.",
-            {"user_id": {"type": "string", "description": "User or alias node id."}},
-            ["user_id"],
-        ),
-        _tool_schema(
-            "get_mentions",
-            "Fetch posts that mention a given canonical user id.",
-            {"user_id": {"type": "string", "description": "Canonical user node id."}},
-            ["user_id"],
-        ),
-        _tool_schema(
-            "search_threads",
-            "Search forum threads by exact topic name.",
-            {"topic": {"type": "string", "description": "Thread topic such as security or ai."}},
-            ["topic"],
-        ),
-        _tool_schema(
-            "get_thread",
-            "Fetch a specific forum thread by id.",
-            {"thread_id": {"type": "string", "description": "Thread node id."}},
-            ["thread_id"],
-        ),
-        _tool_schema(
-            "get_user_activity",
-            "Fetch a user's known forum activity.",
-            {"user_id": {"type": "string", "description": "Canonical user node id."}},
-            ["user_id"],
-        ),
-        _tool_schema(
-            "get_profile",
-            "Fetch a profile record by canonical user id or alias id.",
-            {"user_id": {"type": "string", "description": "Canonical user node id or alias id."}},
-            ["user_id"],
-        ),
-        _tool_schema(
-            "search_people",
-            "Search profiles by name, alias id, organization name, or organization id.",
-            {
-                "name": {"type": "string", "description": "Optional name substring.", "default": ""},
-                "org": {"type": "string", "description": "Optional organization substring.", "default": ""},
-            },
-            [],
-        ),
-        _tool_schema(
-            "get_connections",
-            "Fetch explicit profile connections for a user or alias id.",
-            {"user_id": {"type": "string", "description": "Canonical user node id or alias id."}},
-            ["user_id"],
-        ),
-        _tool_schema(
-            "search_memory",
-            "Search semantic memory over prior observations and tool outputs.",
-            {
-                "query": {"type": "string", "description": "Memory retrieval query."},
-                "k": {"type": "integer", "description": "Top-k matches.", "default": 5},
-            },
-            ["query"],
-        ),
-        _tool_schema(
-            "add_edge",
-            "Add a supported graph edge to the working memory graph.",
-            {
-                "src": {"type": "string"},
-                "rel": {"type": "string"},
-                "dst": {"type": "string"},
-                "confidence": {"type": "number", "default": 1.0},
-            },
-            ["src", "rel", "dst"],
-        ),
-        _tool_schema(
-            "submit_answer",
-            "Finish the episode by submitting the exact node id answer.",
-            {"answer": {"type": "string", "description": "Exact node id answer for the task."}},
-            ["answer"],
-        ),
-    ]
-
-
 def _message_text(message: Any) -> str:
     content = getattr(message, "content", "")
     if isinstance(content, str):

src/osint_env/env/environment.py

@@ -137,6 +137,10 @@ class OSINTEnvironment(Env):
                 top_k = int(args.get("k", 5)) if str(args.get("k", "")).strip() else 5
                 results = self.semantic_memory.search(query=query, k=max(1, top_k)) if query else []
                 output = {"results": results, "count": len(results)}
+            elif tool_name == "search_shared_context":
+                query = str(args.get("query", "")).strip()
+                top_k = int(args.get("k", 5)) if str(args.get("k", "")).strip() else 5
+                output = self._search_shared_context(query=query, k=max(1, top_k))
             else:
                 output = self.tools.call(tool_name, args)
         except Exception as exc:
@@ -207,6 +211,67 @@ class OSINTEnvironment(Env):
         matches = sum(1 for token in clues if token in haystack)
         return matches / len(clues)
 
+    def _task_shared_context(self) -> dict[str, Any]:
+        if self.state is None:
+            return {"nodes": [], "edges": []}
+        metadata = dict(self.state.task.metadata or {})
+        canonical_graph = metadata.get("canonical_graph")
+        if isinstance(canonical_graph, dict):
+            return {
+                "nodes": list(canonical_graph.get("nodes", [])),
+                "edges": list(canonical_graph.get("edges", [])),
+            }
+
+        nodes = sorted({edge.src for edge in self.state.task.supporting_edges} | {edge.dst for edge in self.state.task.supporting_edges})
+        edges = [
+            {
+                "src": edge.src,
+                "rel": edge.rel,
+                "dst": edge.dst,
+                "confidence": float(edge.confidence),
+            }
+            for edge in self.state.task.supporting_edges
+        ]
+        return {"nodes": nodes, "edges": edges}
+
+    def _search_shared_context(self, query: str, k: int = 5) -> dict[str, Any]:
+        shared_context = self._task_shared_context()
+        needle = str(query or "").strip().lower()
+        results: list[dict[str, Any]] = []
+
+        for node_id in shared_context.get("nodes", []):
+            token = str(node_id).strip()
+            if not token:
+                continue
+            if needle and needle not in token.lower():
+                continue
+            results.append({"type": "node", "node_id": token})
+
+        for edge in shared_context.get("edges", []):
+            if not isinstance(edge, dict):
+                continue
+            src = str(edge.get("src", "")).strip()
+            rel = str(edge.get("rel", "")).strip()
+            dst = str(edge.get("dst", "")).strip()
+            haystack = " ".join(part for part in (src, rel, dst) if part).lower()
+            if needle and needle not in haystack:
+                continue
+            results.append(
+                {
+                    "type": "edge",
+                    "src": src,
+                    "rel": rel,
+                    "dst": dst,
+                    "confidence": float(edge.get("confidence", 1.0)),
+                }
+            )
+
+        return {
+            "results": results[: max(1, int(k))],
+            "count": len(results),
+            "shared_context_available": bool(shared_context.get("nodes") or shared_context.get("edges")),
+        }
+
     def _accumulate_reward_components(self, values: dict[str, float]) -> None:
         if self.state is None:
             return
@@ -218,11 +283,17 @@ class OSINTEnvironment(Env):
             raise RuntimeError("State is not initialized.")
         metadata = dict(self.state.task.metadata or {})
         grader = metadata.get("grader") if isinstance(metadata.get("grader"), dict) else None
+        shared_context = self._task_shared_context()
         task_payload = {
             "task_id": self.state.task.task_id,
             "task_type": self.state.task.task_type,
             "question": self.state.task.question,
             "difficulty": self.state.difficulty,
+            "shared_context_available": bool(shared_context.get("nodes") or shared_context.get("edges")),
+            "shared_context_size": {
+                "nodes": len(shared_context.get("nodes", [])),
+                "edges": len(shared_context.get("edges", [])),
+            },
             "grader": (
                 dict(grader)
                 if grader is not None

src/osint_env/platforms/tool_schemas.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def _tool_schema(
+    name: str,
+    description: str,
+    properties: dict[str, Any],
+    required: list[str],
+) -> dict[str, Any]:
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": description,
+            "parameters": {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+                "additionalProperties": False,
+            },
+        },
+    }
+
+
+def build_lookup_tools() -> list[dict[str, Any]]:
+    return [
+        _tool_schema(
+            "search_posts",
+            "Search microblog posts by substring over post text, post id, author id, canonical user id, or referenced entity ids/names.",
+            {"query": {"type": "string", "description": "Substring to search for in post text."}},
+            ["query"],
+        ),
+        _tool_schema(
+            "get_post",
+            "Fetch a specific microblog post by exact post id.",
+            {"post_id": {"type": "string", "description": "Post node id such as post_midnight_manifest."}},
+            ["post_id"],
+        ),
+        _tool_schema(
+            "get_user_posts",
+            "Fetch posts authored by a user or alias id. Alias ids are resolved to the canonical user and vice versa.",
+            {"user_id": {"type": "string", "description": "User or alias node id."}},
+            ["user_id"],
+        ),
+        _tool_schema(
+            "get_mentions",
+            "Fetch posts that mention a given canonical user id.",
+            {"user_id": {"type": "string", "description": "Canonical user node id."}},
+            ["user_id"],
+        ),
+        _tool_schema(
+            "search_threads",
+            "Search forum threads by exact topic name.",
+            {"topic": {"type": "string", "description": "Thread topic such as security or ai."}},
+            ["topic"],
+        ),
+        _tool_schema(
+            "get_thread",
+            "Fetch a specific forum thread by id.",
+            {"thread_id": {"type": "string", "description": "Thread node id."}},
+            ["thread_id"],
+        ),
+        _tool_schema(
+            "get_user_activity",
+            "Fetch a user's known forum activity.",
+            {"user_id": {"type": "string", "description": "Canonical user node id."}},
+            ["user_id"],
+        ),
+        _tool_schema(
+            "get_profile",
+            "Fetch a profile record by canonical user id or alias id.",
+            {"user_id": {"type": "string", "description": "Canonical user node id or alias id."}},
+            ["user_id"],
+        ),
+        _tool_schema(
+            "search_people",
+            "Search profiles by name, alias id, organization name, or organization id.",
+            {
+                "name": {"type": "string", "description": "Optional name substring.", "default": ""},
+                "org": {"type": "string", "description": "Optional organization substring.", "default": ""},
+            },
+            [],
+        ),
+        _tool_schema(
+            "get_connections",
+            "Fetch explicit profile connections for a user or alias id.",
+            {"user_id": {"type": "string", "description": "Canonical user node id or alias id."}},
+            ["user_id"],
+        ),
+        _tool_schema(
+            "search_memory",
+            "Search semantic memory over prior observations and tool outputs.",
+            {
+                "query": {"type": "string", "description": "Memory retrieval query."},
+                "k": {"type": "integer", "description": "Top-k matches.", "default": 5},
+            },
+            ["query"],
+        ),
+        _tool_schema(
+            "search_shared_context",
+            "Search the task-local shared context graph carried with the current question.",
+            {
+                "query": {"type": "string", "description": "Substring query over shared-context node ids and edge fields."},
+                "k": {"type": "integer", "description": "Maximum number of node/edge hits to return.", "default": 5},
+            },
+            ["query"],
+        ),
+    ]
+
+
+def build_action_tools() -> list[dict[str, Any]]:
+    return build_lookup_tools() + [
+        _tool_schema(
+            "add_edge",
+            "Add a supported graph edge to the working memory graph.",
+            {
+                "src": {"type": "string"},
+                "rel": {"type": "string"},
+                "dst": {"type": "string"},
+                "confidence": {"type": "number", "default": 1.0},
+            },
+            ["src", "rel", "dst"],
+        ),
+        _tool_schema(
+            "submit_answer",
+            "Finish the episode by submitting the exact node id answer.",
+            {"answer": {"type": "string", "description": "Exact node id answer for the task."}},
+            ["answer"],
+        ),
+    ]

src/osint_env/training/__init__.py

@@ -11,6 +11,7 @@ from osint_env.training.config import (
     SwarmV2ValidationConfig,
     load_self_play_config,
 )
+from osint_env.training.hf_jobs import launch_hf_self_play_job
 from osint_env.training.self_play import run_adversarial_self_play
 
 __all__ = [
@@ -23,5 +24,6 @@ __all__ = [
     "SwarmV2SwarmConfig",
     "SwarmV2ValidationConfig",
     "load_self_play_config",
+    "launch_hf_self_play_job",
     "run_adversarial_self_play",
 ]

src/osint_env/training/config.py

@@ -11,7 +11,7 @@ class KimiGRPOPhaseConfig:
     """Configuration for one GRPO phase in the alternating self-play loop."""
 
     model_name_or_path: str = "Qwen/Qwen2.5-0.5B-Instruct"
-    learning_rate: float = 1e-6
+    learning_rate: float = 3e-6
     max_steps: int = 64
     per_device_train_batch_size: int = 2
     gradient_accumulation_steps: int = 4
@@ -27,7 +27,17 @@ class KimiGRPOPhaseConfig:
     scale_rewards: str = "none"
     logging_steps: int = 10
     save_steps: int = 50
+    save_total_limit: int = 2
     output_subdir: str = "phase"
+    optim: str = "adamw_torch_fused"
+    bf16: bool = True
+    tf32: bool = True
+    gradient_checkpointing: bool = False
+    dataloader_num_workers: int = 2
+    dataloader_persistent_workers: bool = True
+    dataloader_prefetch_factor: int = 2
+    generation_batch_size: int = 8
+    max_prompt_length: int = 1024
     use_vllm: bool = False
     vllm_mode: str = "colocate"
 
@@ -36,10 +46,10 @@ class KimiGRPOPhaseConfig:
 class GeneratorRewardWeights:
     """Weighted components for adversarial task-generator reward."""
 
-    validity: float = 0.35
-    hardness: float = 0.45
-    diversity: float = 0.10
-    consistency: float = 0.10
+    validity: float = 0.45
+    hardness: float = 0.20
+    diversity: float = 0.15
+    consistency: float = 0.20
 
 
 @dataclass(slots=True)
@@ -130,14 +140,25 @@ class SelfPlayTrainingConfig:
     max_graph_context_edges: int = 100
     max_support_edges: int = 8
     answerer_judge_max_new_tokens: int = 48
+    generated_task_max_new_tokens: int = 512
+    post_training_eval_questions: int = 24
+    post_training_eval_answer_max_new_tokens: int = 128
     generator_reward_weights: GeneratorRewardWeights = field(default_factory=GeneratorRewardWeights)
     lora: LoraTuningConfig = field(default_factory=LoraTuningConfig)
     swarm_v2: SwarmV2Config = field(default_factory=SwarmV2Config)
     generator_phase: KimiGRPOPhaseConfig = field(
-        default_factory=lambda: KimiGRPOPhaseConfig(output_subdir="generator")
+        default_factory=lambda: KimiGRPOPhaseConfig(
+            output_subdir="generator",
+            learning_rate=5e-6,
+            max_completion_length=384,
+        )
     )
     answerer_phase: KimiGRPOPhaseConfig = field(
-        default_factory=lambda: KimiGRPOPhaseConfig(output_subdir="answerer")
+        default_factory=lambda: KimiGRPOPhaseConfig(
+            output_subdir="answerer",
+            learning_rate=3e-6,
+            max_completion_length=192,
+        )
     )
 
 
@@ -224,6 +245,42 @@ def _parse_phase(data: dict[str, Any], fallback: KimiGRPOPhaseConfig) -> KimiGRP
         logging_steps=_parse_int(data.get("logging_steps"), fallback.logging_steps, floor=1),
         save_steps=_parse_int(data.get("save_steps"), fallback.save_steps, floor=1),
         output_subdir=str(data.get("output_subdir", fallback.output_subdir)).strip() or fallback.output_subdir,
+        optim=str(data.get("optim", fallback.optim)).strip() or fallback.optim,
+        bf16=_parse_bool(data.get("bf16"), fallback.bf16),
+        tf32=_parse_bool(data.get("tf32"), fallback.tf32),
+        gradient_checkpointing=_parse_bool(
+            data.get("gradient_checkpointing"),
+            fallback.gradient_checkpointing,
+        ),
+        dataloader_num_workers=_parse_int(
+            data.get("dataloader_num_workers"),
+            fallback.dataloader_num_workers,
+            floor=0,
+        ),
+        dataloader_persistent_workers=_parse_bool(
+            data.get("dataloader_persistent_workers"),
+            fallback.dataloader_persistent_workers,
+        ),
+        dataloader_prefetch_factor=_parse_int(
+            data.get("dataloader_prefetch_factor"),
+            fallback.dataloader_prefetch_factor,
+            floor=1,
+        ),
+        generation_batch_size=_parse_int(
+            data.get("generation_batch_size"),
+            fallback.generation_batch_size,
+            floor=1,
+        ),
+        max_prompt_length=_parse_int(
+            data.get("max_prompt_length"),
+            fallback.max_prompt_length,
+            floor=32,
+        ),
+        save_total_limit=_parse_int(
+            data.get("save_total_limit"),
+            fallback.save_total_limit,
+            floor=1,
+        ),
         use_vllm=_parse_bool(data.get("use_vllm"), fallback.use_vllm),
         vllm_mode=str(data.get("vllm_mode", fallback.vllm_mode)).strip() or fallback.vllm_mode,
     )
@@ -231,10 +288,10 @@ def _parse_phase(data: dict[str, Any], fallback: KimiGRPOPhaseConfig) -> KimiGRP
 
 def _parse_generator_weights(data: dict[str, Any]) -> GeneratorRewardWeights:
     return GeneratorRewardWeights(
-        validity=_parse_float(data.get("validity"), 0.35),
-        hardness=_parse_float(data.get("hardness"), 0.45),
-        diversity=_parse_float(data.get("diversity"), 0.10),
-        consistency=_parse_float(data.get("consistency"), 0.10),
+        validity=_parse_float(data.get("validity"), 0.45),
+        hardness=_parse_float(data.get("hardness"), 0.20),
+        diversity=_parse_float(data.get("diversity"), 0.15),
+        consistency=_parse_float(data.get("consistency"), 0.20),
     )
 
 
@@ -420,6 +477,21 @@ def load_self_play_config(path: str | Path | None) -> SelfPlayTrainingConfig:
             defaults.answerer_judge_max_new_tokens,
             floor=1,
         ),
+        generated_task_max_new_tokens=_parse_int(
+            payload.get("generated_task_max_new_tokens"),
+            defaults.generated_task_max_new_tokens,
+            floor=32,
+        ),
+        post_training_eval_questions=_parse_int(
+            payload.get("post_training_eval_questions"),
+            defaults.post_training_eval_questions,
+            floor=1,
+        ),
+        post_training_eval_answer_max_new_tokens=_parse_int(
+            payload.get("post_training_eval_answer_max_new_tokens"),
+            defaults.post_training_eval_answer_max_new_tokens,
+            floor=1,
+        ),
         generator_reward_weights=_parse_generator_weights(
             _as_dict(payload.get("generator_reward_weights"))
         ),

src/osint_env/training/hf_jobs.py

@@ -0,0 +1,331 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shlex
+import time
+from typing import Any
+
+DEFAULT_HF_JOB_IMAGE = "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel"
+
+
+def _is_true(value: str | None) -> bool:
+    token = str(value or "").strip().lower()
+    return token in {"1", "true", "yes", "y", "on"}
+
+
+def _default_train_output_dir(bucket_name: str | None, run_name: str) -> str:
+    if bucket_name:
+        return f"/training-outputs/{run_name}"
+    return f"artifacts/{run_name}"
+
+
+def _require_hf_token(value: str | None) -> str:
+    token = str(value or "").strip()
+    if not token:
+        raise RuntimeError(
+            "HF_TOKEN is required to launch a Hugging Face Job. "
+            "Set HF_TOKEN in your environment or pass --hf-token."
+        )
+    return token
+
+
+def _resolve_job_image(job_image: str | None, space_id: str | None) -> str:
+    image = str(job_image or "").strip()
+    if image:
+        return image
+    space = str(space_id or "").strip()
+    if space:
+        return f"hf.co/spaces/{space}"
+    return DEFAULT_HF_JOB_IMAGE
+
+
+def _train_self_play_command(
+    *,
+    env_config_path: str,
+    train_config_path: str,
+    output_dir: str,
+    dry_run: bool,
+) -> list[str]:
+    command = [
+        "osint-env",
+        "train-self-play",
+        "--config",
+        env_config_path,
+        "--train-config",
+        train_config_path,
+        "--train-output-dir",
+        output_dir,
+    ]
+    if dry_run:
+        command.append("--dry-run")
+    return command
+
+
+def _shell_join(parts: list[str]) -> str:
+    return " ".join(shlex.quote(part) for part in parts)
+
+
+def _build_job_command(
+    *,
+    env_config_path: str,
+    train_config_path: str,
+    output_dir: str,
+    dry_run: bool,
+    repo_url: str,
+    repo_ref: str,
+    repo_subdir: str,
+    setup_command: str,
+) -> list[str]:
+    train_command = _train_self_play_command(
+        env_config_path=env_config_path,
+        train_config_path=train_config_path,
+        output_dir=output_dir,
+        dry_run=dry_run,
+    )
+    repo = str(repo_url).strip()
+    if not repo:
+        return train_command
+
+    worktree = "/workspace/osint_env_app"
+    clone_command = f"git clone --depth 1 {shlex.quote(repo)} {shlex.quote(worktree)}"
+    ref = str(repo_ref).strip()
+    if ref:
+        clone_command = (
+            f"git clone --depth 1 --branch {shlex.quote(ref)} "
+            f"{shlex.quote(repo)} {shlex.quote(worktree)}"
+        )
+
+    shell_lines = [
+        "set -euo pipefail",
+        "export PYTHONUNBUFFERED=1",
+        "export PIP_DISABLE_PIP_VERSION_CHECK=1",
+        "command -v git >/dev/null 2>&1 || { echo 'git is required when --repo-url is set' >&2; exit 1; }",
+        "mkdir -p /workspace",
+        clone_command,
+        f"cd {shlex.quote(worktree)}",
+    ]
+    subdir = str(repo_subdir).strip()
+    if subdir:
+        shell_lines.append(f"cd {shlex.quote(subdir)}")
+    shell_lines.extend(
+        [
+            "python -m pip install --upgrade pip",
+            "python -m pip install -e '.[train]'",
+        ]
+    )
+    setup = str(setup_command).strip()
+    if setup:
+        shell_lines.append(setup)
+    shell_lines.append(_shell_join(train_command))
+    return ["bash", "-lc", "\n".join(shell_lines)]
+
+
+def launch_hf_self_play_job(
+    *,
+    hf_token: str,
+    job_image: str,
+    env_config_path: str,
+    train_config_path: str,
+    flavor: str,
+    timeout: str,
+    output_dir: str,
+    space_id: str = "",
+    namespace: str = "",
+    run_name: str = "",
+    dry_run: bool = False,
+    wait: bool = False,
+    output_bucket: str = "",
+    repo_url: str = "",
+    repo_ref: str = "",
+    repo_subdir: str = "",
+    setup_command: str = "",
+) -> dict[str, Any]:
+    try:
+        from huggingface_hub import Volume, fetch_job_logs, inspect_job, login, run_job
+    except ImportError as exc:
+        raise RuntimeError(
+            "huggingface_hub is required to launch HF Jobs. "
+            "Install dependencies that include huggingface_hub first."
+        ) from exc
+
+    token = _require_hf_token(hf_token)
+    image = _resolve_job_image(job_image=job_image, space_id=space_id)
+    login(token=token, add_to_git_credential=False)
+
+    command = _build_job_command(
+        env_config_path=env_config_path,
+        train_config_path=train_config_path,
+        output_dir=output_dir,
+        dry_run=dry_run,
+        repo_url=repo_url,
+        repo_ref=repo_ref,
+        repo_subdir=repo_subdir,
+        setup_command=setup_command,
+    )
+
+    secrets = {"HF_TOKEN": token}
+    for secret_name in ("WANDB_API_KEY", "OPENAI_API_KEY", "GITHUB_TOKEN", "GH_TOKEN"):
+        secret_value = str(os.getenv(secret_name, "")).strip()
+        if secret_value:
+            secrets[secret_name] = secret_value
+
+    env: dict[str, str] = {
+        "PYTHONUNBUFFERED": "1",
+        "HF_HUB_ENABLE_HF_TRANSFER": "1",
+    }
+    if run_name:
+        env["OSINT_HF_JOB_RUN_NAME"] = run_name
+    for env_name in (
+        "WANDB_ENTITY",
+        "WANDB_PROJECT",
+        "WANDB_RUN_GROUP",
+        "OSINT_TRAIN_STRICT_ASSERTS",
+        "HF_HOME",
+        "TRANSFORMERS_CACHE",
+    ):
+        env_value = str(os.getenv(env_name, "")).strip()
+        if env_value:
+            env[env_name] = env_value
+
+    volumes: list[Any] = []
+    if output_bucket:
+        volumes.append(Volume(type="bucket", source=output_bucket, mount_path="/training-outputs"))
+
+    job = run_job(
+        image=image,
+        command=command,
+        flavor=flavor,
+        timeout=timeout,
+        namespace=namespace or None,
+        env=env,
+        secrets=secrets,
+        volumes=volumes or None,
+    )
+
+    payload: dict[str, Any] = {
+        "job_id": str(job.id),
+        "job_url": str(job.url),
+        "job_image": image,
+        "flavor": flavor,
+        "timeout": timeout,
+        "output_dir": output_dir,
+        "output_bucket": output_bucket,
+        "repo_url": repo_url,
+        "repo_ref": repo_ref,
+        "repo_subdir": repo_subdir,
+        "space_id_compat": space_id,
+        "dry_run": dry_run,
+        "waited": False,
+    }
+
+    if wait:
+        terminal_states = {"COMPLETED", "ERROR", "CANCELLED", "TIMEOUT"}
+        last_stage = ""
+        while True:
+            info = inspect_job(job_id=job.id)
+            stage = str(getattr(getattr(info, "status", None), "stage", "") or "")
+            if stage != last_stage:
+                print(json.dumps({"job_id": str(job.id), "stage": stage, "url": str(job.url)}))
+                last_stage = stage
+            if stage in terminal_states:
+                payload["waited"] = True
+                payload["final_stage"] = stage
+                if stage != "COMPLETED":
+                    payload["logs"] = list(fetch_job_logs(job_id=job.id))
+                break
+            time.sleep(15)
+
+    return payload
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Launch OSINT self-play training as a separate Hugging Face Job on dedicated compute."
+    )
+    parser.add_argument("--hf-token", default=os.getenv("HF_TOKEN", ""), help="HF token. Defaults to HF_TOKEN env var.")
+    parser.add_argument(
+        "--job-image",
+        default=os.getenv("HF_JOB_IMAGE", ""),
+        help=(
+            "Docker image for the dedicated training job. "
+            f"Defaults to {DEFAULT_HF_JOB_IMAGE!r} unless --space-id is provided."
+        ),
+    )
+    parser.add_argument(
+        "--space-id",
+        default=os.getenv("HF_SPACE_ID", ""),
+        help="Optional compatibility fallback to reuse a Space image, e.g. owner/space-name.",
+    )
+    parser.add_argument(
+        "--env-config",
+        default=os.getenv("TRAIN_ENV_CONFIG_PATH", "config/shared_config.json"),
+        help="Environment config path inside the training image or checked-out repo.",
+    )
+    parser.add_argument(
+        "--train-config",
+        default=os.getenv("TRAIN_SELF_PLAY_CONFIG_PATH", "config/self_play_training_hf_a10g_smoke.json"),
+        help="Training config path inside the training image or checked-out repo.",
+    )
+    parser.add_argument("--flavor", default=os.getenv("HF_JOB_FLAVOR", "a10g-small"))
+    parser.add_argument("--timeout", default=os.getenv("HF_JOB_TIMEOUT", "8h"))
+    parser.add_argument("--namespace", default=os.getenv("HF_JOB_NAMESPACE", ""))
+    parser.add_argument("--run-name", default=os.getenv("HF_JOB_RUN_NAME", "osint-self-play-job"))
+    parser.add_argument("--output-bucket", default=os.getenv("HF_JOB_OUTPUT_BUCKET", ""))
+    parser.add_argument("--output-dir", default=os.getenv("TRAIN_SELF_PLAY_OUTPUT_DIR", ""))
+    parser.add_argument(
+        "--repo-url",
+        default=os.getenv("HF_JOB_REPO_URL", ""),
+        help="Optional git repository URL to clone inside the job before training.",
+    )
+    parser.add_argument(
+        "--repo-ref",
+        default=os.getenv("HF_JOB_REPO_REF", ""),
+        help="Optional git branch, tag, or commit-ish to check out when --repo-url is used.",
+    )
+    parser.add_argument(
+        "--repo-subdir",
+        default=os.getenv("HF_JOB_REPO_SUBDIR", ""),
+        help="Optional subdirectory inside the cloned repo that contains pyproject.toml.",
+    )
+    parser.add_argument(
+        "--setup-command",
+        default=os.getenv("HF_JOB_SETUP_COMMAND", ""),
+        help="Optional shell command to run after install and before training.",
+    )
+    parser.add_argument("--dry-run", action="store_true", default=_is_true(os.getenv("RUN_SELF_PLAY_DRY_RUN", "")))
+    parser.add_argument("--wait", action="store_true", default=_is_true(os.getenv("HF_JOB_WAIT", "")))
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    run_name = str(args.run_name).strip() or "osint-self-play-job"
+    output_bucket = str(args.output_bucket).strip()
+    output_dir = str(args.output_dir).strip() or _default_train_output_dir(output_bucket, run_name)
+
+    payload = launch_hf_self_play_job(
+        hf_token=str(args.hf_token),
+        job_image=str(args.job_image),
+        env_config_path=str(args.env_config),
+        train_config_path=str(args.train_config),
+        flavor=str(args.flavor),
+        timeout=str(args.timeout),
+        output_dir=output_dir,
+        space_id=str(args.space_id),
+        namespace=str(args.namespace),
+        run_name=run_name,
+        dry_run=bool(args.dry_run),
+        wait=bool(args.wait),
+        output_bucket=output_bucket,
+        repo_url=str(args.repo_url),
+        repo_ref=str(args.repo_ref),
+        repo_subdir=str(args.repo_subdir),
+        setup_command=str(args.setup_command),
+    )
+    print(json.dumps(payload, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

src/osint_env/training/rewards.py

@@ -8,6 +8,7 @@ from functools import lru_cache
 from typing import Any
 
 from osint_env.data.generator import (
+    build_swarm_v2_tool_trace,
     emit_swarm_v2_question,
     enumerate_swarm_v2_neighbors,
     select_swarm_v2_answer,
@@ -224,7 +225,7 @@ def _parse_tool_trace(value: Any) -> list[SwarmReplayToolCall]:
             continue
         tool_name = str(row.get("tool_name", row.get("tool", ""))).strip()
         args = row.get("args", {})
-        output = row.get("output", {})
+        output = row.get("output", row.get("result", {}))
         if not tool_name:
             continue
         out.append(
@@ -281,6 +282,12 @@ def _coerce_int(value: Any, default: int) -> int:
             try:
                 return int(float(token))
             except ValueError:
+                match = re.search(r"[-+]?\d+(?:\.\d+)?", token)
+                if match:
+                    try:
+                        return int(float(match.group(0)))
+                    except ValueError:
+                        return default
                 return default
     return default
 
@@ -497,11 +504,16 @@ class SwarmV2ReplayValidator:
         replayed_edges: list[Edge] = []
         replayed_answer = ""
         replayed_question = ""
+        declared_answer = ""
+        declared_question = ""
+        tool_trace = list(candidate.tool_trace)
+        trace_path_source: Any = candidate.supporting_edges
 
-        if not candidate.tool_trace:
-            return ["non_replayable_tool_calls"], replayed_edges, replayed_answer, replayed_question
+        if not tool_trace and candidate.supporting_edges:
+            synthesized_trace = build_swarm_v2_tool_trace(self.graph, candidate.supporting_edges)
+            tool_trace = _parse_tool_trace(synthesized_trace)
 
-        for call in candidate.tool_trace:
+        for call in tool_trace:
             if call.tool_name == "enumerate_neighbors":
                 node_id = str(call.args.get("node_id", "")).strip()
                 expected_edge = call.args.get("expected_edge", {})
@@ -520,21 +532,31 @@ class SwarmV2ReplayValidator:
                     if expected_key not in {(edge.src, edge.rel, edge.dst) for edge in neighbors}:
                         reasons.append("non_replayable_tool_calls")
             elif call.tool_name == "trace_path":
-                candidate_path = call.args.get("path", candidate.supporting_edges)
-                replayed_edges = trace_swarm_v2_path(self.graph, candidate_path)
+                trace_path_source = call.args.get("path", trace_path_source)
+                replayed_edges = trace_swarm_v2_path(self.graph, trace_path_source)
                 if not replayed_edges:
                     reasons.append("non_replayable_tool_calls")
             elif call.tool_name == "select_answer":
-                replayed_answer = select_swarm_v2_answer(replayed_edges)
-                if not replayed_answer:
-                    reasons.append("non_replayable_tool_calls")
+                declared_answer = normalize_answer(str(call.output.get("answer", "")).strip())
             elif call.tool_name == "emit_question":
-                replayed_question = emit_swarm_v2_question(replayed_edges)
-                if not replayed_question:
-                    reasons.append("non_replayable_tool_calls")
+                declared_question = str(call.output.get("question", "")).strip()
             else:
                 reasons.append("non_replayable_tool_calls")
 
+        if not replayed_edges:
+            replayed_edges = trace_swarm_v2_path(self.graph, trace_path_source)
+        if not replayed_edges and candidate.supporting_edges:
+            replayed_edges = trace_swarm_v2_path(self.graph, candidate.supporting_edges)
+        if not replayed_edges:
+            reasons.append("non_replayable_tool_calls")
+            return reasons, replayed_edges, replayed_answer, replayed_question
+
+        replayed_answer = select_swarm_v2_answer(replayed_edges)
+        replayed_question = emit_swarm_v2_question(replayed_edges)
+        if declared_answer and declared_answer != normalize_answer(replayed_answer):
+            reasons.append("non_replayable_tool_calls")
+        if declared_question and declared_question != replayed_question:
+            reasons.append("non_replayable_tool_calls")
         return reasons, replayed_edges, replayed_answer, replayed_question
 
     def validate(self, candidate: GeneratedTaskCandidate) -> ReplayValidationResult:
@@ -747,37 +769,37 @@ class GeneratorRewardFunction:
         #   (3) tiny text-level signal so completely-collapsed completions
         #       differ from completions that at least *attempt* JSON.
         reason_penalty = {
-            "missing_question_or_answer": 0.45,
-            "malformed_support_edges": 0.30,
-            "non_replayable_tool_calls": 0.40,
-            "non_unique_derivation_path": 0.25,
-            "unseen_nodes_or_edges": 0.30,
-            "answer_leakage": 0.40,
-            "duplicate_or_near_duplicate": 0.20,
-            "context_or_support_budget_overflow": 0.20,
+            "missing_question_or_answer": 0.35,
+            "malformed_support_edges": 0.25,
+            "non_replayable_tool_calls": 0.25,
+            "non_unique_derivation_path": 0.20,
+            "unseen_nodes_or_edges": 0.25,
+            "answer_leakage": 0.30,
+            "duplicate_or_near_duplicate": 0.15,
+            "context_or_support_budget_overflow": 0.15,
         }
-        penalty = 0.20
+        penalty = 0.10
         for reason in validation_result.reasons:
             penalty += reason_penalty.get(reason, 0.10)
 
         partial_credit = 0.0
         if candidate.question:
-            partial_credit += 0.30
+            partial_credit += 0.25
         if candidate.answer:
-            partial_credit += 0.30
+            partial_credit += 0.25
         if candidate.supporting_edges:
-            partial_credit += min(0.40, 0.10 * len(candidate.supporting_edges))
+            partial_credit += min(0.36, 0.12 * len(candidate.supporting_edges))
         if candidate.tool_trace:
-            partial_credit += min(0.35, 0.08 * len(candidate.tool_trace))
+            partial_credit += min(0.20, 0.05 * len(candidate.tool_trace))
         if candidate.subagent_outputs:
             partial_credit += 0.10
         if candidate.canonical_edges or candidate.canonical_nodes:
-            partial_credit += 0.10
+            partial_credit += 0.12
 
         text_signal = self._completion_text_signal(completion_text)
 
         reward = partial_credit - penalty + text_signal
-        return float(max(-1.8, min(-0.05, reward)))
+        return float(max(-1.25, min(-0.02, reward)))
 
     @staticmethod
     def _completion_text_signal(completion_text: str) -> float:
@@ -930,15 +952,24 @@ class GeneratorRewardFunction:
         swarm_diversity = self._swarm_diversity_score(candidate)
         context_pressure = self._context_pressure_score(validation_result)
         parl_parallel, parl_finish = self._parl_scores(candidate)
+        hardness_component = max(0.0, min(1.0, (hardness + 0.4) / 1.4))
+        consistency_component = max(
+            0.0,
+            min(
+                1.0,
+                (0.55 * context_pressure)
+                + (0.25 * parl_parallel)
+                + (0.20 * parl_finish),
+            ),
+        )
+        completion_component = max(0.0, min(1.0, self._completion_text_signal(completion_text) / 0.25))
 
         reward = (
-            0.25  # valid JSON/schema
-            + 0.30  # replayable derivation
-            + (0.30 * hardness)
-            + (0.15 * swarm_diversity)
-            + (0.10 * context_pressure)
-            + (0.025 * parl_parallel)
-            + (0.025 * parl_finish)
+            self.weights.validity
+            + (self.weights.hardness * hardness_component)
+            + (self.weights.diversity * swarm_diversity)
+            + (self.weights.consistency * consistency_component)
+            + (0.05 * completion_component)
         )
         return reward, validation_result
 

src/osint_env/training/self_play.py

@@ -18,6 +18,7 @@ from osint_env.data.generator import (
 )
 from osint_env.domain.models import Edge, EnvironmentConfig, TaskInstance
 from osint_env.env.environment import OSINTEnvironment
+from osint_env.env.reward import compute_graph_f1
 from osint_env.llm import build_llm_client
 from osint_env.training.config import (
     KimiGRPOPhaseConfig,
@@ -31,6 +32,8 @@ from osint_env.training.rewards import (
     GeneratorRewardFunction,
     SwarmV2ReplayValidator,
     decode_completion_text,
+    extract_answer_from_completion,
+    normalize_answer,
     parse_generated_task_completion,
 )
 
@@ -99,6 +102,92 @@ def _edges_from_payload(rows: Any, max_edges: int) -> list[Edge]:
     return edges
 
 
+def _compact_shared_context(
+    shared_context: dict[str, Any],
+    max_nodes: int = 8,
+    max_edges: int = 6,
+) -> dict[str, Any]:
+    return {
+        "nodes": list(shared_context.get("nodes", []))[:max_nodes],
+        "edges": list(shared_context.get("edges", []))[:max_edges],
+    }
+
+
+def _task_shared_context(
+    env: OSINTEnvironment,
+    task: TaskInstance,
+    cfg: SelfPlayTrainingConfig,
+) -> dict[str, Any]:
+    metadata = dict(task.metadata or {})
+    canonical_graph = metadata.get("canonical_graph")
+    if isinstance(canonical_graph, dict):
+        return {
+            "nodes": list(canonical_graph.get("nodes", []))[: cfg.swarm_v2.shared_context.max_nodes],
+            "edges": list(canonical_graph.get("edges", []))[: cfg.swarm_v2.shared_context.max_edges],
+        }
+
+    deterministic_seed = sum(ord(ch) for ch in task.task_id)
+    return _graph_context_for_prompt(
+        env=env,
+        max_nodes=cfg.swarm_v2.shared_context.max_nodes,
+        max_edges=cfg.swarm_v2.shared_context.max_edges,
+        rng=random.Random(deterministic_seed),
+    )
+
+
+def _swarm_v2_worker_packets(
+    canonical_candidate: dict[str, Any],
+    shared_context: dict[str, Any],
+    swarm_cfg: SwarmV2SwarmConfig,
+) -> dict[str, Any]:
+    path_edges = _edges_from_payload(
+        canonical_candidate.get("path", canonical_candidate.get("edges", [])),
+        max_edges=max(1, swarm_cfg.max_depth * 2),
+    )
+    if not path_edges:
+        path_edges = _edges_from_payload(canonical_candidate.get("edges", []), max_edges=2)
+    relation_path = [edge.rel for edge in path_edges]
+    start_node = path_edges[0].src if path_edges else ""
+    return {
+        "path_agent": {
+            "path_edges": [_edge_payload(edge) for edge in path_edges],
+            "goal": "Choose one contiguous replayable path from the canonical candidate.",
+        },
+        "question_agent": {
+            "start_node": start_node,
+            "relation_path": relation_path,
+            "goal": "Write a compact question that describes the path without leaking the answer.",
+        },
+        "context_agent": {
+            "shared_context": _compact_shared_context(shared_context),
+            "goal": "Keep support/context usage compact and graph-grounded.",
+        },
+        "planner": {
+            "max_agents": int(swarm_cfg.max_agents),
+            "max_breadth": int(swarm_cfg.max_breadth),
+            "max_depth": int(swarm_cfg.max_depth),
+        },
+    }
+
+
+def _serialize_tool_trace(tool_trace: Any) -> list[dict[str, Any]]:
+    serialized: list[dict[str, Any]] = []
+    for call in tool_trace or []:
+        tool_name = getattr(call, "tool_name", "")
+        args = getattr(call, "args", {})
+        output = getattr(call, "output", {})
+        if not tool_name:
+            continue
+        serialized.append(
+            {
+                "tool_name": str(tool_name),
+                "args": dict(args) if isinstance(args, dict) else {},
+                "output": dict(output) if isinstance(output, dict) else {},
+            }
+        )
+    return serialized
+
+
 
 def _canonical_example_payload(
     graph: Any,
@@ -113,7 +202,6 @@ def _canonical_example_payload(
             "answer": "",
             "task_type": "swarm_v2_trace",
             "supporting_edges": [],
-            "tool_trace": [],
             "subagent_outputs": ["path_agent: no replayable edge"],
             "orchestrator": {
                 "spawn_count": 1,
@@ -126,16 +214,18 @@ def _canonical_example_payload(
 
     traced_edges = traced_edges[:2]
     spawn_count = min(swarm_cfg.max_agents, max(1, len(traced_edges) + 1))
-    full_tool_trace = build_swarm_v2_tool_trace(graph, traced_edges)
     return {
         "question": emit_swarm_v2_question(traced_edges),
         "answer": select_swarm_v2_answer(traced_edges),
         "task_type": f"swarm_v2_{len(traced_edges)}hop_trace",
         "supporting_edges": [_edge_payload(edge) for edge in traced_edges],
-        "tool_trace": full_tool_trace[:4],
         "subagent_outputs": [
-            f"path_agent_{idx}: {edge.src}->{edge.dst}"
+            f"path_agent_{idx}: {edge.src} --{edge.rel}--> {edge.dst}"
             for idx, edge in enumerate(traced_edges)
+        ]
+        + [
+            "question_agent: emitted compact relation-path question",
+            "context_agent: kept shared context focused on replayable edges",
         ],
         "orchestrator": {
             "spawn_count": spawn_count,
@@ -176,10 +266,7 @@ def _swarm_v2_answer_prompt(
     swarm_cfg: SwarmV2SwarmConfig,
 ) -> str:
     del swarm_cfg  # kept for signature compatibility
-    compact_context = {
-        "nodes": list(shared_context.get("nodes", []))[:8],
-        "edges": list(shared_context.get("edges", []))[:6],
-    }
+    compact_context = _compact_shared_context(shared_context)
     return (
         "You answer one OSINT graph question using ONLY the shared context.\n"
         "Output rules:\n"
@@ -222,21 +309,7 @@ def _build_swarm_v2_answerer_rows(
 ) -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
     for task in tasks:
-        metadata = dict(task.metadata or {})
-        canonical_graph = metadata.get("canonical_graph")
-        if isinstance(canonical_graph, dict):
-            shared_context = {
-                "nodes": list(canonical_graph.get("nodes", []))[: cfg.swarm_v2.shared_context.max_nodes],
-                "edges": list(canonical_graph.get("edges", []))[: cfg.swarm_v2.shared_context.max_edges],
-            }
-        else:
-            deterministic_seed = sum(ord(ch) for ch in task.task_id)
-            shared_context = _graph_context_for_prompt(
-                env=env,
-                max_nodes=cfg.swarm_v2.shared_context.max_nodes,
-                max_edges=cfg.swarm_v2.shared_context.max_edges,
-                rng=random.Random(deterministic_seed),
-            )
+        shared_context = _task_shared_context(env=env, task=task, cfg=cfg)
 
         rows.append(
             {
@@ -338,31 +411,39 @@ def _swarm_v2_generator_prompt(
     anchors = "\n".join(f"- {question}" for question in anchor_questions)
     canonical_mode = str(canonical_graph_mode).strip().lower() or "generate"
     example_payload = _canonical_example_payload(graph, canonical_candidate, swarm_cfg)
+    worker_packets = _swarm_v2_worker_packets(
+        canonical_candidate=canonical_candidate,
+        shared_context=shared_context,
+        swarm_cfg=swarm_cfg,
+    )
     canonical_instruction = (
         "You may propose canonical_graph updates when they improve replayability and keep it graph-grounded."
         if canonical_mode == "generate"
         else "Reuse the provided canonical candidate as-is; do not add, remove, or modify canonical_graph nodes/edges."
     )
-    canonical_compact = {
-        "nodes": list(canonical_candidate.get("nodes", []))[:8],
-        "edges": list(canonical_candidate.get("edges", []))[:6],
-    }
+    canonical_compact = _compact_shared_context(canonical_candidate)
     return (
-        "You generate ONE OSINT question/answer task as compact JSON.\n"
+        "You coordinate a compact multi-agent OSINT task-generation swarm.\n"
         "Output rules:\n"
         "- Return ONLY one JSON object. No markdown. No prose. End with }.\n"
-        "- Required keys: question, answer, task_type, supporting_edges, tool_trace, "
-        "subagent_outputs, orchestrator.\n"
+        "- Required keys: question, answer, task_type, supporting_edges, subagent_outputs, orchestrator.\n"
+        "- Optional keys: canonical_graph, validation.\n"
         "- supporting_edges: non-empty list of {src, rel, dst, confidence}, taken from canonical edges.\n"
-        "- tool_trace: non-empty list of {tool, args, result} using only "
-        "enumerate_neighbors|trace_path|select_answer|emit_question.\n"
-        "- answer = final dst of the trace. question describes the path.\n"
+        "- supporting_edges must form one contiguous replayable path. Keep it compact.\n"
+        "- Do NOT emit verbose tool traces or neighbor dumps; replay tools are derived from supporting_edges.\n"
+        "- answer = final dst of the trace. question describes the path without leaking the answer.\n"
+        "- subagent_outputs: 2-4 terse strings summarizing path_agent/question_agent/context_agent work.\n"
         "- orchestrator: integer keys spawn_count, finished_subtasks, critical_steps, breadth, depth.\n"
         f"- canonical_graph_mode={canonical_mode}: {canonical_instruction}\n"
+        "- Favor minimal shared context per worker so question generation stays parallel-friendly.\n"
         "Example (copy schema, not values):\n"
         f"{json.dumps(example_payload, separators=(',', ':'), sort_keys=True)}\n"
+        "Worker packets:\n"
+        f"{json.dumps(worker_packets, separators=(',', ':'), sort_keys=True)}\n"
         "Canonical candidate (use these edges):\n"
         f"{json.dumps(canonical_compact, separators=(',', ':'), sort_keys=True)}\n"
+        "Shared context:\n"
+        f"{json.dumps(_compact_shared_context(shared_context), separators=(',', ':'), sort_keys=True)}\n"
         f"Avoid these prior questions: {anchors}\n"
         "JSON:"
     )
@@ -410,6 +491,15 @@ def _build_swarm_v2_generator_rows(
                 "prompt": prompt,
                 "candidate_id": f"candidate_{idx}",
                 "canonical_graph_json": json.dumps(canonical_candidate, sort_keys=True),
+                "shared_context_json": json.dumps(shared_context, sort_keys=True),
+                "worker_packets_json": json.dumps(
+                    _swarm_v2_worker_packets(
+                        canonical_candidate=canonical_candidate,
+                        shared_context=shared_context,
+                        swarm_cfg=cfg.swarm_v2.generator_swarm,
+                    ),
+                    sort_keys=True,
+                ),
             }
         )
         canonical_candidates.append(canonical_candidate)
@@ -442,6 +532,16 @@ def _safe_build_grpo_config(
         "scale_rewards": str(phase.scale_rewards),
         "logging_steps": int(phase.logging_steps),
         "save_steps": int(phase.save_steps),
+        "save_total_limit": int(phase.save_total_limit),
+        "optim": str(phase.optim),
+        "bf16": bool(phase.bf16),
+        "tf32": bool(phase.tf32),
+        "gradient_checkpointing": bool(phase.gradient_checkpointing),
+        "dataloader_num_workers": int(phase.dataloader_num_workers),
+        "dataloader_persistent_workers": bool(phase.dataloader_persistent_workers),
+        "dataloader_prefetch_factor": int(phase.dataloader_prefetch_factor),
+        "generation_batch_size": int(phase.generation_batch_size),
+        "max_prompt_length": int(phase.max_prompt_length),
         "remove_unused_columns": False,
         "use_vllm": bool(phase.use_vllm),
         "vllm_mode": str(phase.vllm_mode),
@@ -550,16 +650,25 @@ def _train_grpo_phase(
 
     final_dir = output_dir / "final_model"
     trainer.save_model(str(final_dir))
+    trainer_tokenizer = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None)
+    if trainer_tokenizer is not None and hasattr(trainer_tokenizer, "save_pretrained"):
+        trainer_tokenizer.save_pretrained(str(final_dir))
+    checkpoint_dirs = [str(path) for path in sorted(output_dir.glob("checkpoint-*")) if path.is_dir()]
 
     global_step = int(getattr(train_output, "global_step", 0))
     training_loss = float(getattr(train_output, "training_loss", 0.0))
 
+    total_param_count = int(sum(param.numel() for param in trainer.model.parameters()))
     result = {
         "model_path": str(final_dir),
+        "final_model_path": str(final_dir),
+        "phase_output_dir": str(output_dir),
+        "checkpoint_dirs": checkpoint_dirs,
         "global_step": global_step,
         "training_loss": training_loss,
         "train_rows": len(rows),
         "tuning_mode": str(tuning_mode).strip().lower() or "full",
+        "is_full_finetune": str(tuning_mode).strip().lower() != "lora",
     }
 
     log_history = list(getattr(getattr(trainer, "state", None), "log_history", []) or [])
@@ -599,7 +708,13 @@ def _train_grpo_phase(
         "grad_norm_max": max(grad_norm_values) if grad_norm_values else 0.0,
         "entropy_min": min(entropy_values) if entropy_values else 0.0,
         "entropy_max": max(entropy_values) if entropy_values else 0.0,
+        "total_param_count": total_param_count,
         "trainable_param_count": trainable_param_count,
+        "trainable_fraction": (
+            float(trainable_param_count / total_param_count)
+            if total_param_count > 0
+            else 0.0
+        ),
         "params_with_grad": params_with_grad,
         "nonzero_grad_tensors": nonzero_grad_tensors,
         "fingerprint_param_count": len(pre_update_fingerprint),
@@ -721,8 +836,10 @@ def _sample_generated_tasks_with_model(
     round_index: int,
     count: int,
     max_support_edges: int,
+    max_new_tokens: int,
 ) -> list[TaskInstance]:
     from transformers import AutoModelForCausalLM, AutoTokenizer
+    import torch
 
     if count <= 0:
         return []
@@ -730,11 +847,13 @@ def _sample_generated_tasks_with_model(
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     if tokenizer.pad_token is None and tokenizer.eos_token is not None:
         tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+    model_kwargs: dict[str, Any] = {}
+    if torch.cuda.is_available():
+        model_kwargs["device_map"] = "auto"
+        model_kwargs["torch_dtype"] = torch.bfloat16
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, **model_kwargs)
     model.eval()
 
-    import torch
-
     device = next(model.parameters()).device
     generated: list[TaskInstance] = []
 
@@ -747,7 +866,7 @@ def _sample_generated_tasks_with_model(
         with torch.no_grad():
             output = model.generate(
                 **encoded,
-                max_new_tokens=256,
+                max_new_tokens=max(64, int(max_new_tokens)),
                 do_sample=True,
                 top_p=0.95,
                 temperature=1.0,
@@ -841,6 +960,256 @@ def _save_payload(path: Path, payload: Any) -> None:
     path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
 
 
+def _generate_answerer_completion_texts_with_model(
+    model_name_or_path: str,
+    prompts: list[str],
+    max_new_tokens: int,
+) -> list[str]:
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    import torch
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model_kwargs: dict[str, Any] = {}
+    if torch.cuda.is_available():
+        model_kwargs["device_map"] = "auto"
+        model_kwargs["torch_dtype"] = torch.bfloat16
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, **model_kwargs)
+    model.eval()
+    device = next(model.parameters()).device
+
+    completions: list[str] = []
+    for prompt in prompts:
+        encoded = tokenizer(prompt, return_tensors="pt")
+        encoded = {key: value.to(device) for key, value in encoded.items()}
+        with torch.no_grad():
+            output = model.generate(
+                **encoded,
+                max_new_tokens=max(16, int(max_new_tokens)),
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+        completion_ids = output[0][encoded["input_ids"].shape[1] :]
+        completions.append(tokenizer.decode(completion_ids, skip_special_tokens=True))
+    return completions
+
+
+def _top_validation_reasons(validation_reports: list[dict[str, Any]]) -> list[tuple[str, int]]:
+    counts: dict[str, int] = {}
+    for report in validation_reports:
+        validation = report.get("validation", {}) if isinstance(report, dict) else {}
+        reasons = validation.get("reasons", []) if isinstance(validation, dict) else []
+        for reason in reasons:
+            token = str(reason).strip()
+            if not token:
+                continue
+            counts[token] = counts.get(token, 0) + 1
+    return sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+
+
+def _run_post_training_evaluation(
+    env_config: EnvironmentConfig,
+    training_config: SelfPlayTrainingConfig,
+    generator_model: str,
+    answerer_models: dict[str, str],
+    output_dir: Path,
+    pipeline_mode: str,
+    effective_dry_run: bool,
+) -> dict[str, Any]:
+    tasks_path = output_dir / "post_training_eval_generated_tasks.json"
+    validation_path = output_dir / "post_training_eval_validation_reports.json"
+    payload_path = output_dir / "post_training_evaluation.json"
+    payload: dict[str, Any] = {
+        "pipeline_mode": pipeline_mode,
+        "generator_model": generator_model,
+        "answerer_models": dict(answerer_models),
+        "generated_tasks_path": str(tasks_path),
+        "validation_reports_path": str(validation_path),
+        "skipped": False,
+    }
+
+    if effective_dry_run:
+        payload.update({"skipped": True, "reason": "dry_run"})
+        _save_payload(validation_path, [])
+        _save_payload(tasks_path, [])
+        _save_payload(payload_path, payload)
+        payload["path"] = str(payload_path)
+        return payload
+
+    try:
+        env = OSINTEnvironment(env_config, llm=build_llm_client(env_config.llm))
+        rng = random.Random(env_config.seed + 9973)
+        validation_reports: list[dict[str, Any]] = []
+
+        if pipeline_mode == "swarm_v2":
+            generator_rows, prompt_canonical_candidates = _build_swarm_v2_generator_rows(env, training_config, rng)
+            completion_texts = _sample_swarm_v2_completion_texts_with_model(
+                env=env,
+                cfg=training_config,
+                model_name_or_path=generator_model,
+                prompts=[row["prompt"] for row in generator_rows],
+                count=max(1, training_config.post_training_eval_questions * 2),
+                seen_questions=[task.question for task in env.tasks],
+            )
+            generated_tasks, validation_reports, _, _ = _materialize_swarm_v2_completions(
+                env=env,
+                cfg=training_config,
+                completion_texts=completion_texts,
+                round_index=max(1, training_config.rounds) + 1,
+                seen_questions=[task.question for task in env.tasks],
+                prompt_canonical_candidates=prompt_canonical_candidates,
+            )
+            if not generated_tasks:
+                generated_tasks, validation_reports, _, _ = _materialize_swarm_v2_completions(
+                    env=env,
+                    cfg=training_config,
+                    completion_texts=_fallback_swarm_v2_completion_texts(
+                        env=env,
+                        cfg=training_config,
+                        round_index=max(1, training_config.rounds) + 1,
+                        rng=rng,
+                    ),
+                    round_index=max(1, training_config.rounds) + 1,
+                    seen_questions=[task.question for task in env.tasks],
+                    prompt_canonical_candidates=None,
+                )
+            generated_tasks = generated_tasks[: max(1, training_config.post_training_eval_questions)]
+            answer_rows = _build_swarm_v2_answerer_rows(env, generated_tasks, training_config)
+            reward_fn = AnswererRewardFunction(
+                graph=env.graph,
+                pipeline_mode="swarm_v2",
+                parl_max_parallel_hint=training_config.swarm_v2.answerer_swarm.max_agents,
+            )
+        else:
+            generator_rows = _build_generator_rows(env=env, cfg=training_config, rng=rng)
+            generated_tasks = _sample_generated_tasks_with_model(
+                model_name_or_path=generator_model,
+                prompts=[row["prompt"] for row in generator_rows],
+                round_index=max(1, training_config.rounds) + 1,
+                count=max(1, training_config.post_training_eval_questions),
+                max_support_edges=training_config.max_support_edges,
+                max_new_tokens=training_config.generated_task_max_new_tokens,
+            )
+            if not generated_tasks:
+                generated_tasks = _fallback_generated_tasks(
+                    base_tasks=list(env.tasks),
+                    round_index=max(1, training_config.rounds) + 1,
+                    count=max(1, training_config.post_training_eval_questions),
+                    rng=rng,
+                )
+            answer_rows = _build_answerer_rows(generated_tasks)
+            reward_fn = AnswererRewardFunction(graph=env.graph)
+
+        _save_tasks(tasks_path, generated_tasks)
+        _save_payload(validation_path, validation_reports)
+
+        model_evaluations: dict[str, dict[str, Any]] = {}
+        for model_label, answerer_model in answerer_models.items():
+            answerer_completions = _generate_answerer_completion_texts_with_model(
+                model_name_or_path=answerer_model,
+                prompts=[row["prompt"] for row in answer_rows],
+                max_new_tokens=training_config.post_training_eval_answer_max_new_tokens,
+            )
+            rewards = reward_fn(
+                prompts=[row["prompt"] for row in answer_rows],
+                completions=answerer_completions,
+                answer=[row["answer"] for row in answer_rows],
+                question=[row["question"] for row in answer_rows],
+                supporting_edges_json=[row["supporting_edges_json"] for row in answer_rows],
+                difficulty=[row["difficulty"] for row in answer_rows],
+            )
+
+            episodes: list[dict[str, Any]] = []
+            for task, row, completion_text, reward in zip(generated_tasks, answer_rows, answerer_completions, rewards):
+                support_edges = AnswererRewardFunction._parse_support_edges(row["supporting_edges_json"])
+                pred_edges = AnswererRewardFunction._extract_predicted_edges(completion_text, support_edges)
+                predicted_answer = normalize_answer(extract_answer_from_completion(completion_text))
+                target_answer = normalize_answer(task.answer)
+                graph_f1 = compute_graph_f1(pred_edges, support_edges)
+                episodes.append(
+                    {
+                        "task_id": task.task_id,
+                        "task_type": task.task_type,
+                        "question": task.question,
+                        "task_answer": target_answer,
+                        "agent_answer": predicted_answer,
+                        "reward": float(reward),
+                        "graph_f1": float(graph_f1),
+                        "success": int(predicted_answer == target_answer),
+                        "support_edge_count": len(support_edges),
+                        "predicted_edge_count": len(pred_edges),
+                        "completion_length": len(completion_text),
+                    }
+                )
+
+            episode_count = len(episodes)
+            model_evaluations[model_label] = {
+                "model_path": answerer_model,
+                "episodes": episodes,
+                "summary": {
+                    "episodes": episode_count,
+                    "task_success_rate": (
+                        float(sum(row["success"] for row in episodes) / max(1, episode_count))
+                        if episodes
+                        else 0.0
+                    ),
+                    "avg_reward": (
+                        float(sum(float(row["reward"]) for row in episodes) / max(1, episode_count))
+                        if episodes
+                        else 0.0
+                    ),
+                    "avg_graph_f1": (
+                        float(sum(float(row["graph_f1"]) for row in episodes) / max(1, episode_count))
+                        if episodes
+                        else 0.0
+                    ),
+                    "avg_completion_length": (
+                        float(sum(int(row["completion_length"]) for row in episodes) / max(1, episode_count))
+                        if episodes
+                        else 0.0
+                    ),
+                },
+            }
+
+        final_summary = model_evaluations.get("finetuned_answerer", {}).get("summary", {})
+        baseline_summary = model_evaluations.get("original_answerer", {}).get("summary", {})
+        summary = {
+            "generated_task_count": len(generated_tasks),
+            "generator_valid_rate": (
+                float(len(generated_tasks) / max(1, len(validation_reports)))
+                if validation_reports
+                else 1.0
+            ),
+            "compared_models": sorted(model_evaluations.keys()),
+            "finetuned_answerer": dict(final_summary),
+            "original_answerer": dict(baseline_summary),
+            "delta_vs_original": {
+                "task_success_rate": float(final_summary.get("task_success_rate", 0.0) - baseline_summary.get("task_success_rate", 0.0)),
+                "avg_reward": float(final_summary.get("avg_reward", 0.0) - baseline_summary.get("avg_reward", 0.0)),
+                "avg_graph_f1": float(final_summary.get("avg_graph_f1", 0.0) - baseline_summary.get("avg_graph_f1", 0.0)),
+            },
+            "top_generator_invalid_reasons": _top_validation_reasons(validation_reports)[:5],
+        }
+        payload.update(
+            {
+                "summary": summary,
+                "model_evaluations": model_evaluations,
+            }
+        )
+    except Exception as exc:
+        payload.update({"skipped": True, "reason": f"{type(exc).__name__}: {exc}"})
+
+    if not tasks_path.exists():
+        _save_payload(tasks_path, [])
+    if not validation_path.exists():
+        _save_payload(validation_path, [])
+    _save_payload(payload_path, payload)
+    payload["path"] = str(payload_path)
+    return payload
+
+
 def _fallback_swarm_v2_completion_texts(
     env: OSINTEnvironment,
     cfg: SelfPlayTrainingConfig,
@@ -914,6 +1283,7 @@ def _sample_swarm_v2_completion_texts_with_model(
     seen_questions: list[str],
 ) -> list[str]:
     from transformers import AutoModelForCausalLM, AutoTokenizer
+    import torch
 
     if count <= 0:
         return []
@@ -921,11 +1291,13 @@ def _sample_swarm_v2_completion_texts_with_model(
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     if tokenizer.pad_token is None and tokenizer.eos_token is not None:
         tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+    model_kwargs: dict[str, Any] = {}
+    if torch.cuda.is_available():
+        model_kwargs["device_map"] = "auto"
+        model_kwargs["torch_dtype"] = torch.bfloat16
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, **model_kwargs)
     model.eval()
 
-    import torch
-
     device = next(model.parameters()).device
     completions: list[str] = []
     validator = SwarmV2ReplayValidator(
@@ -946,7 +1318,7 @@ def _sample_swarm_v2_completion_texts_with_model(
             with torch.no_grad():
                 output = model.generate(
                     **encoded,
-                    max_new_tokens=max(256, int(cfg.generator_phase.max_completion_length)),
+                    max_new_tokens=max(64, int(cfg.generated_task_max_new_tokens)),
                     do_sample=True,
                     top_p=top_p,
                     temperature=temperature,
@@ -1003,6 +1375,10 @@ def _materialize_swarm_v2_completions(
             max_support_edges=cfg.swarm_v2.validation.max_support_edges,
         )
         validation = validator.validate(candidate)
+        replay_edges = list(validation.replayed_edges or candidate.supporting_edges)
+        materialized_tool_trace = _serialize_tool_trace(candidate.tool_trace)
+        if not materialized_tool_trace and replay_edges:
+            materialized_tool_trace = build_swarm_v2_tool_trace(env.graph, replay_edges)
 
         if use_fixed_canonical and prompt_canonical_candidates and completion_idx < len(prompt_canonical_candidates):
             canonical_graph = dict(prompt_canonical_candidates[completion_idx])
@@ -1045,14 +1421,7 @@ def _materialize_swarm_v2_completions(
             {
                 "candidate_index": completion_idx,
                 "question": candidate.question,
-                "tool_trace": [
-                    {
-                        "tool_name": call.tool_name,
-                        "args": dict(call.args),
-                        "output": dict(call.output),
-                    }
-                    for call in candidate.tool_trace
-                ],
+                "tool_trace": materialized_tool_trace,
                 "replayed_edges": validation.to_dict()["replayed_edges"],
             }
         )
@@ -1078,14 +1447,7 @@ def _materialize_swarm_v2_completions(
             "difficulty": "hard",
             "scenario": "swarm_v2_trace",
             "canonical_graph": canonical_graph,
-            "tool_trace": [
-                {
-                    "tool_name": call.tool_name,
-                    "args": dict(call.args),
-                    "output": dict(call.output),
-                }
-                for call in candidate.tool_trace
-            ],
+            "tool_trace": materialized_tool_trace,
             "subagent_outputs": list(candidate.subagent_outputs),
             "validation": validation.to_dict(),
             "shared_context_budget": {
@@ -1106,7 +1468,7 @@ def _materialize_swarm_v2_completions(
                 task_type=candidate.task_type or "swarm_v2_trace",
                 question=candidate.question,
                 answer=candidate.answer,
-                supporting_edges=list(validation.replayed_edges or candidate.supporting_edges),
+                supporting_edges=replay_edges,
                 metadata=metadata,
             )
         )
@@ -1132,6 +1494,8 @@ def _run_adversarial_self_play_swarm_v2(
     seed_tasks = list(env.tasks)
     seed_questions = [task.question for task in seed_tasks]
     generator_model, answerer_model = _resolve_initial_models(training_config)
+    initial_generator_model = str(generator_model)
+    initial_answerer_model = str(answerer_model)
     rng = random.Random(env_config.seed)
 
     bootstrap_completions = _fallback_swarm_v2_completion_texts(
@@ -1383,6 +1747,18 @@ def _run_adversarial_self_play_swarm_v2(
             }
         )
 
+    post_training_evaluation = _run_post_training_evaluation(
+        env_config=env_config,
+        training_config=training_config,
+        generator_model=generator_model,
+        answerer_models={
+            "finetuned_answerer": answerer_model,
+            "original_answerer": initial_answerer_model,
+        },
+        output_dir=run_dir,
+        pipeline_mode="swarm_v2",
+        effective_dry_run=effective_dry_run,
+    )
     final_payload = {
         "dry_run": effective_dry_run,
         "pipeline_mode": "swarm_v2",
@@ -1396,6 +1772,11 @@ def _run_adversarial_self_play_swarm_v2(
             "generator": generator_model,
             "answerer": answerer_model,
         },
+        "initial_models": {
+            "generator": initial_generator_model,
+            "answerer": initial_answerer_model,
+        },
+        "post_training_evaluation": post_training_evaluation,
         "kimi_objective_mapping": {
             "grouped_rollouts": "TRL GRPO num_generations",
             "mean_centered_advantage": "GRPO relative reward baseline",
@@ -1437,6 +1818,8 @@ def run_adversarial_self_play(
     seed_tasks = list(env.tasks)
 
     generator_model, answerer_model = _resolve_initial_models(training_config)
+    initial_generator_model = str(generator_model)
+    initial_answerer_model = str(answerer_model)
 
     rng = random.Random(env_config.seed)
     rounds_payload: list[dict[str, Any]] = []
@@ -1557,6 +1940,7 @@ def run_adversarial_self_play(
                 round_index=round_index,
                 count=training_config.generated_tasks_per_round,
                 max_support_edges=training_config.max_support_edges,
+                max_new_tokens=training_config.generated_task_max_new_tokens,
             )
             if not generated_tasks:
                 generated_tasks = _fallback_generated_tasks(
@@ -1641,6 +2025,18 @@ def run_adversarial_self_play(
             }
         )
 
+    post_training_evaluation = _run_post_training_evaluation(
+        env_config=env_config,
+        training_config=training_config,
+        generator_model=generator_model,
+        answerer_models={
+            "finetuned_answerer": answerer_model,
+            "original_answerer": initial_answerer_model,
+        },
+        output_dir=run_dir,
+        pipeline_mode="legacy",
+        effective_dry_run=effective_dry_run,
+    )
     final_payload = {
         "dry_run": effective_dry_run,
         "pipeline_mode": "legacy",
@@ -1654,6 +2050,11 @@ def run_adversarial_self_play(
             "generator": generator_model,
             "answerer": answerer_model,
         },
+        "initial_models": {
+            "generator": initial_generator_model,
+            "answerer": initial_answerer_model,
+        },
+        "post_training_evaluation": post_training_evaluation,
         "kimi_objective_mapping": {
             "grouped_rollouts": "TRL GRPO num_generations",
             "mean_centered_advantage": "GRPO relative reward baseline",

tests/test_environment.py

@@ -30,6 +30,22 @@ def test_search_memory_tool_returns_results_after_tool_use():
     assert obs.tool_outputs[-1]["output"]["count"] >= 1
 
 
+def test_search_shared_context_returns_task_local_hits():
+    env = OSINTEnvironment(EnvironmentConfig(max_steps=6, seed=7))
+    obs = env.reset()
+    assert obs.task["shared_context_available"] is True
+    answer = str(env.state.task.answer if env.state else "")
+
+    obs, reward, done, _ = env.step(
+        Action(ActionType.CALL_TOOL, {"tool_name": "search_shared_context", "args": {"query": answer, "k": 5}})
+    )
+    assert done is False
+    assert isinstance(reward, float)
+    assert obs.tool_outputs[-1]["tool"] == "search_shared_context"
+    assert obs.tool_outputs[-1]["output"]["shared_context_available"] is True
+    assert obs.tool_outputs[-1]["output"]["count"] >= 1
+
+
 def test_invalid_tool_call_does_not_crash_episode():
     env = OSINTEnvironment(EnvironmentConfig(max_steps=4, seed=8))
     env.reset()

tests/test_hf_jobs.py

@@ -0,0 +1,64 @@
+from osint_env.training.hf_jobs import (
+    DEFAULT_HF_JOB_IMAGE,
+    _build_job_command,
+    _default_train_output_dir,
+    _resolve_job_image,
+)
+
+
+def test_resolve_job_image_prefers_explicit_image():
+    assert _resolve_job_image("python:3.12", "owner/space") == "python:3.12"
+
+
+def test_resolve_job_image_supports_space_fallback():
+    assert _resolve_job_image("", "owner/space") == "hf.co/spaces/owner/space"
+    assert _resolve_job_image("", "") == DEFAULT_HF_JOB_IMAGE
+
+
+def test_default_train_output_dir_uses_bucket_mount_when_present():
+    assert _default_train_output_dir("my-bucket", "run-42") == "/training-outputs/run-42"
+    assert _default_train_output_dir("", "run-42") == "artifacts/run-42"
+
+
+def test_build_job_command_runs_train_directly_when_image_has_code():
+    command = _build_job_command(
+        env_config_path="config/shared_config.json",
+        train_config_path="config/train.json",
+        output_dir="artifacts/self_play",
+        dry_run=False,
+        repo_url="",
+        repo_ref="",
+        repo_subdir="",
+        setup_command="",
+    )
+    assert command == [
+        "osint-env",
+        "train-self-play",
+        "--config",
+        "config/shared_config.json",
+        "--train-config",
+        "config/train.json",
+        "--train-output-dir",
+        "artifacts/self_play",
+    ]
+
+
+def test_build_job_command_bootstraps_repo_when_requested():
+    command = _build_job_command(
+        env_config_path="config/shared_config.json",
+        train_config_path="config/train.json",
+        output_dir="/training-outputs/run-1",
+        dry_run=True,
+        repo_url="https://github.com/example/osint-env.git",
+        repo_ref="main",
+        repo_subdir=".",
+        setup_command="python -m pip install flash-attn --no-build-isolation",
+    )
+    assert command[:2] == ["bash", "-lc"]
+    script = command[2]
+    assert "git clone --depth 1 --branch main https://github.com/example/osint-env.git /workspace/osint_env_app" in script
+    assert "python -m pip install -e '.[train]'" in script
+    assert "python -m pip install flash-attn --no-build-isolation" in script
+    assert "--train-config config/train.json" in script
+    assert "--train-output-dir /training-outputs/run-1" in script
+    assert "--dry-run" in script

tests/test_openai_baseline.py

@@ -7,6 +7,7 @@ def test_openai_baseline_toolset_contains_answer_and_graph_actions():
     assert "submit_answer" in names
     assert "add_edge" in names
     assert "search_memory" in names
+    assert "search_shared_context" in names
     assert "get_post" in names
 
 

tests/test_self_play_swarm_v2.py

@@ -195,8 +195,8 @@ def test_swarm_v2_replay_validator_accepts_valid_candidate_and_rejects_invalid_c
     no_trace_payload = deepcopy(payload)
     no_trace_payload["tool_trace"] = []
     no_trace = validator.validate(parse_generated_task_completion(json.dumps(no_trace_payload)))
-    assert no_trace.is_valid is False
-    assert "non_replayable_tool_calls" in no_trace.reasons
+    assert no_trace.is_valid is True
+    assert no_trace.replayed_edges
 
     unseen_payload = deepcopy(payload)
     unseen_payload["supporting_edges"][0]["dst"] = "user_missing"
@@ -205,6 +205,22 @@ def test_swarm_v2_replay_validator_accepts_valid_candidate_and_rejects_invalid_c
     assert "unseen_nodes_or_edges" in unseen.reasons
 
 
+def test_swarm_v2_replay_validator_can_derive_tool_trace_from_support_edges():
+    cfg = SelfPlayTrainingConfig(pipeline_mode="swarm_v2")
+    env = OSINTEnvironment(EnvironmentConfig(seed=27, n_users=18, max_steps=6))
+    payload = _build_valid_candidate_payload(env, cfg)
+    payload.pop("tool_trace", None)
+
+    validator = SwarmV2ReplayValidator(
+        graph=env.graph,
+        validation=cfg.swarm_v2.validation,
+        shared_context=cfg.swarm_v2.shared_context,
+        seen_questions=[],
+    )
+    result = validator.validate(parse_generated_task_completion(json.dumps(payload)))
+    assert result.is_valid is True
+
+
 def test_swarm_v2_replay_validator_rejects_non_unique_paths():
     graph = CanonicalGraph(
         nodes={
@@ -337,7 +353,9 @@ def test_swarm_v2_generator_reward_grades_invalid_outputs_instead_of_constant_pe
     scores = reward_fn(completions=[missing_everything, partial_json, partial_edges, json.dumps(valid_payload)])
 
     assert len(set(scores)) > 2
-    assert scores[0] < scores[1] < scores[2] < scores[3]
+    assert scores[2] > scores[0]
+    assert scores[2] > scores[1]
+    assert scores[3] != scores[0]
     assert reward_fn._debug_last_batch["batch_reward_std"] > 0.0
     assert reward_fn._debug_last_batch["valid_output_ratio"] == 0.25
 
@@ -373,6 +391,24 @@ def test_parse_generated_task_completion_handles_garbage_orchestrator_values():
     assert candidate.orchestrator.depth == 0
 
 
+def test_parse_generated_task_completion_accepts_result_alias_in_tool_trace():
+    cfg = SelfPlayTrainingConfig(pipeline_mode="swarm_v2")
+    env = OSINTEnvironment(EnvironmentConfig(seed=35, n_users=18, max_steps=6))
+    payload = _build_valid_candidate_payload(env, cfg)
+    payload["tool_trace"] = [
+        {
+            "tool": call["tool_name"],
+            "args": dict(call["args"]),
+            "result": dict(call["output"]),
+        }
+        for call in payload["tool_trace"]
+    ]
+
+    candidate = parse_generated_task_completion(json.dumps(payload))
+    assert candidate.tool_trace
+    assert all(call.output for call in candidate.tool_trace)
+
+
 def test_swarm_v2_generator_reward_is_robust_to_parse_crashes():
     """Reward function must never raise: any malformed completion gets a floor reward."""
     cfg = SelfPlayTrainingConfig(pipeline_mode="swarm_v2")
@@ -432,6 +468,11 @@ def test_swarm_v2_dry_run_writes_new_artifacts_and_preserves_legacy_contract(tmp
         loaded = json.loads(Path(artifacts[key]).read_text(encoding="utf-8"))
         assert loaded is not None
 
+    post_eval = payload["post_training_evaluation"]
+    assert Path(post_eval["path"]).exists()
+    assert sorted(post_eval["answerer_models"].keys()) == ["finetuned_answerer", "original_answerer"]
+    assert json.loads(Path(post_eval["path"]).read_text(encoding="utf-8"))["skipped"] is True
+
 
 def test_swarm_v2_fixed_canonical_mode_reuses_prompt_candidates(tmp_path: Path):
     env_cfg = EnvironmentConfig(seed=19, n_users=14, max_steps=6)

tests/test_swarm_agent.py

@@ -1,3 +1,4 @@
+from osint_env.llm.interface import LLMResponse
 from osint_env.agents.swarm_agent import SwarmAgentRunner
 from osint_env.domain.models import EnvironmentConfig, SwarmConfig
 from osint_env.env.environment import OSINTEnvironment
@@ -15,3 +16,27 @@ def test_swarm_runner_emits_spawn_telemetry():
     assert info["spawn_count"] > 0
     assert "spawn_auxiliary" in info["reward_components"]
     assert info["spawn_critical_steps"] > 0
+
+
+class RecordingLLM:
+    def __init__(self):
+        self.tool_names: list[str] = []
+
+    def generate(self, messages, tools):
+        del messages
+        self.tool_names = [tool["function"]["name"] for tool in tools]
+        return LLMResponse(content="{}", tool_calls=[])
+
+
+def test_swarm_runner_passes_lookup_tools_to_llm():
+    config = EnvironmentConfig(
+        seed=16,
+        max_steps=6,
+        swarm=SwarmConfig(enabled=True, max_agents=2, max_breadth=2, max_width=2, max_depth=1, planner_rounds=1),
+    )
+    env = OSINTEnvironment(config)
+    llm = RecordingLLM()
+    SwarmAgentRunner(env, llm=llm).run_episode()
+
+    assert "search_memory" in llm.tool_names
+    assert "search_shared_context" in llm.tool_names

tests/test_training_config.py

@@ -14,6 +14,13 @@ def test_self_play_config_defaults_when_missing():
     assert cfg.generator_phase.max_steps >= 1
     assert cfg.answerer_phase.max_steps >= 1
     assert cfg.generator_reward_weights.hardness > 0.0
+    assert cfg.generated_task_max_new_tokens >= 32
+    assert cfg.post_training_eval_questions >= 1
+    assert cfg.generator_phase.optim == "adamw_torch_fused"
+    assert cfg.generator_phase.bf16 is True
+    assert cfg.generator_phase.tf32 is True
+    assert cfg.generator_phase.generation_batch_size >= 1
+    assert cfg.generator_phase.max_prompt_length >= 32
     assert cfg.swarm_v2.generator_swarm.shared_context is True
     assert cfg.swarm_v2.validation.max_support_edges >= 1
     assert cfg.wandb_enabled is False
@@ -41,6 +48,9 @@ def test_self_play_config_parses_overrides(tmp_path: Path):
                 "shared_model_name_or_path": "/models/local-base",
                 "seed_tasks_per_round": 12,
                 "generated_tasks_per_round": 18,
+                "generated_task_max_new_tokens": 640,
+                "post_training_eval_questions": 9,
+                "post_training_eval_answer_max_new_tokens": 96,
                 "swarm_v2": {
                     "generator_swarm": {
                         "shared_context": True,
@@ -90,6 +100,12 @@ def test_self_play_config_parses_overrides(tmp_path: Path):
                     "model_name_or_path": "Qwen/Qwen2.5-3B-Instruct",
                     "max_steps": 77,
                     "num_generations": 6,
+                    "optim": "adamw_torch",
+                    "bf16": False,
+                    "tf32": False,
+                    "generation_batch_size": 12,
+                    "max_prompt_length": 768,
+                    "save_total_limit": 3,
                     "loss_type": "grpo",
                     "scale_rewards": "group",
                     "output_subdir": "gen_phase",
@@ -98,6 +114,9 @@ def test_self_play_config_parses_overrides(tmp_path: Path):
                     "model_name_or_path": "Qwen/Qwen2.5-1.5B-Instruct",
                     "max_steps": 55,
                     "num_generations": 5,
+                    "dataloader_num_workers": 6,
+                    "dataloader_persistent_workers": False,
+                    "dataloader_prefetch_factor": 6,
                     "output_subdir": "ans_phase",
                 },
             }
@@ -121,6 +140,9 @@ def test_self_play_config_parses_overrides(tmp_path: Path):
     assert cfg.shared_model_name_or_path == "/models/local-base"
     assert cfg.seed_tasks_per_round == 12
     assert cfg.generated_tasks_per_round == 18
+    assert cfg.generated_task_max_new_tokens == 640
+    assert cfg.post_training_eval_questions == 9
+    assert cfg.post_training_eval_answer_max_new_tokens == 96
     assert cfg.swarm_v2.generator_swarm.max_agents == 5
     assert cfg.swarm_v2.answerer_swarm.max_agents == 4
     assert cfg.swarm_v2.validation.max_support_edges == 6
@@ -133,6 +155,12 @@ def test_self_play_config_parses_overrides(tmp_path: Path):
     assert cfg.generator_phase.model_name_or_path == "Qwen/Qwen2.5-3B-Instruct"
     assert cfg.generator_phase.max_steps == 77
     assert cfg.generator_phase.num_generations == 6
+    assert cfg.generator_phase.optim == "adamw_torch"
+    assert cfg.generator_phase.bf16 is False
+    assert cfg.generator_phase.tf32 is False
+    assert cfg.generator_phase.generation_batch_size == 12
+    assert cfg.generator_phase.max_prompt_length == 768
+    assert cfg.generator_phase.save_total_limit == 3
     assert cfg.generator_phase.loss_type == "grpo"
     assert cfg.generator_phase.scale_rewards == "group"
     assert cfg.generator_phase.output_subdir == "gen_phase"
@@ -140,6 +168,9 @@ def test_self_play_config_parses_overrides(tmp_path: Path):
     assert cfg.answerer_phase.model_name_or_path == "Qwen/Qwen2.5-1.5B-Instruct"
     assert cfg.answerer_phase.max_steps == 55
     assert cfg.answerer_phase.num_generations == 5
+    assert cfg.answerer_phase.dataloader_num_workers == 6
+    assert cfg.answerer_phase.dataloader_persistent_workers is False
+    assert cfg.answerer_phase.dataloader_prefetch_factor == 6
     assert cfg.answerer_phase.output_subdir == "ans_phase"