Add constraint-aware safety pipeline: DAgger, safe Pi0 policy, geometry, and scripts

Introduces safety_geometry.py (workspace constraint primitives), safe_pi0_policy.py
(safety-filtered Pi0 wrapper with Python 3.14 argparse shim), dagger.py (DAgger
training loop), and supporting scripts for rollout recording and training.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: jpothen8

AGENTS.md

@@ -181,10 +181,46 @@ uv run python -m sim.scripts.collect_demos --repo-id local/safe-cube-bc \
 .venv/bin/mjpython -m sim.scripts.view_env
 ```
 
-## What's NOT done yet (for the next agent)
+## Training pipeline (`SafePI0Policy` + DAgger) — built, NOT yet run
+
+The policy/training half of `project_summary.md` now exists but is **untested**
+(no runnable Python env was available when it was written — validate before
+trusting it):
+
+- **`sim/safety_geometry.py`** — pure-torch, LeRobot-free primitives: `box_sdf`,
+  `sdf_clearance`, `safety_loss`, and `FKChain` (differentiable FK via
+  `pytorch_kinematics`, end link `gripper_frame_link`). Unit-tested in
+  `sim/tests/test_safety_geometry.py` (`uv run pytest sim/tests -q`). **Run these
+  first** — a sign error here is cheap to catch now, expensive 40 epochs into a 3B run.
+- **`sim/safe_pi0_policy.py`** — `SafePI0Config` (registered as `safe_pi0`) +
+  `SafePI0Policy(PI0Policy)`. Overrides `forward` so one velocity prediction feeds
+  both `L_flow` and the endpoint-estimate safety term. **Flow convention:** LeRobot's
+  π0 is `x_t = t·noise + (1−t)·actions` (t=0→data, t=1→noise) — the *opposite* of
+  `project_summary.md` §5.1. So the endpoint estimate is `â = x_t − t·v` and the
+  safety loss is weighted by `(1−t)`. Actions are MEAN_STD-normalized before
+  `forward`, so `â` is un-normalized (via `dataset_stats[ACTION]`, captured in
+  `__init__`) before FK.
+- **`sim/dagger.py` / `sim/scripts/dagger.py`** — `PolicyRollout` (checkpoint →
+  closed-loop `act(obs)` via the stock preprocessor→`select_action`→postprocessor
+  path) + the collect-round logic, and a CLI orchestrator that runs
+  collect→retrain rounds (retrain shells out to the stock entrypoint).
+- **`sim/scripts/train_safe_pi0.py`** — thin shim: imports the policy module (to
+  register `safe_pi0` with draccus) then calls the stock `train()`.
+- **Privileged-key plumbing (one stock edit):** `src/lerobot/processor/converters.py`
+  `_extract_complementary_data` now preserves any `privileged.*` batch key. Without
+  this the preprocessor (`batch = preprocessor(batch)` in the train loop) drops them
+  before `forward` and the safety loss can't see the cube positions. They ride
+  through as complementary data → un-normalized, untouched by feature classification
+  (`dataset_to_policy_features` `continue`s on non-`observation`/`action` keys).
+
+**Before the first real run, verify:** (1) `pytorch_kinematics` is installed and
+`FKChain.fk(state[:5])` matches the recorded `privileged.ee_pos` on a sample frame
+(FK-frame == world-frame assumption); (2) privileged keys actually arrive in
+`policy.forward`'s batch; (3) sweep `λ` (`--policy.safety_weight`) — it's the
+central knob.
+
+## Still NOT done
 
-- **`SafePI0Policy`** — the LeRobot subclass with `L = L_flow_matching + λ·L_safety` per `project_summary.md` §5. Privileged keys are already in the recorded datasets (`privileged.cube_positions`, `privileged.cube_half_extents`, `privileged.blue_cube_pos`, `privileged.goal_pos`, `privileged.ee_pos`, `privileged.grasped`) — the policy/loss just needs to read them from the batch and ignore them at inference. Per-cube arrays are flattened to 1D; reshape on load.
-- **DAgger driver** — sketched in `sim/README.md`. The primitives it needs are already in place (`SafeCubeEnv`, `ScriptedExpert`, `EpisodeRecorder`, `evaluate.evaluate()`).
 - **Domain randomization not yet applied.** `sim/randomize.py` defines `apply_dr(model, rng, cfg)` but `env.reset()` doesn't call it. Wire it in after the scene compiles, before the first physics step, when ready.
 - **Real-world fine-tune / hardware deployment** — out of scope for the sim.
 

sim/README.md

@@ -98,36 +98,111 @@ from the batch.
 }
 ```
 
-## DAgger loop sketch
+## End-to-end runbook (collect → BC warm-start → DAgger)
 
-The pieces above are the primitives — the DAgger loop is a small driver:
+The policy + training code now exists (`sim/safe_pi0_policy.py`, `sim/dagger.py`,
+`sim/scripts/`). It was written without a runnable environment, so **it is
+untested** — do step 0 before anything else.
+
+### 0. Setup + validate the geometry (cheap, do it first)
+
+```bash
+uv sync --locked --extra all          # base + policy deps
+uv pip install pytorch_kinematics      # differentiable FK for the safety loss
+
+# A sign error in the SDF/FK is cheap to catch now, expensive 40 epochs into a
+# 3B-param fine-tune. These tests are torch-only (no MuJoCo / VLA needed).
+uv run pytest sim/tests -q
+```
+
+Then spot-check the FK frame against ground truth on one recorded frame: load a
+sample, run `FKChain.fk(state[:5])`, and confirm it matches that frame's
+`privileged.ee_pos` (the loss assumes FK's base_link frame == the world frame —
+true while the URDF base sits at the origin).
+
+### 1. Collect data
+
+The safety term needs **both** success and failure data (failures are its
+negative class). Make a clean set for the BC warm start and a larger mixed set
+that contains failures:
+
+```bash
+# Clean successes — BC warm start
+uv run python -m sim.scripts.collect_demos \
+    --repo-id local/safe-cube-bc --root data/safe_cube_bc \
+    --n-episodes 100 --successes-only
+
+# Mixed set (keeps red-contact / drop / timeout failures)
+uv run python -m sim.scripts.collect_demos \
+    --repo-id local/safe-cube-mixed --root data/safe_cube_mixed \
+    --n-episodes 200
+```
+
+### 2. BC warm start (fine-tune π0 with the safety loss on)
+
+```bash
+uv run python -m sim.scripts.train_safe_pi0 \
+    --policy.type=safe_pi0 \
+    --policy.pretrained_path=lerobot/pi0_base \
+    --policy.safety_weight=1.0 \
+    --dataset.repo_id=local/safe-cube-mixed --dataset.root=data/safe_cube_mixed \
+    --output_dir=outputs/safe_pi0_bc \
+    --batch_size=8 --steps=20000 \
+    --policy.gradient_checkpointing=true
+```
+
+`--policy.safety_weight` (λ) is the central knob — sweep it. Too small → safety
+ignored; too large → the policy collapses to pure avoidance and stops doing the
+task. The evolving safety term makes the loss landscape non-stationary, so
+**early-stop at peak rollout success, not at min loss.** Checkpoint lands at
+`outputs/safe_pi0_bc/checkpoints/last/pretrained_model`.
+
+### 3. DAgger (collect on the policy's own state distribution → retrain)
+
+`sim/scripts/dagger.py` runs the full collect→retrain loop. Round 0 is pure
+expert; later rounds load the previous checkpoint, mix in policy actions
+(`alpha = alpha0 ** round`), and **always label with the expert** — which is what
+drags training onto deployment-distribution states and produces fresh failures.
+
+```bash
+uv run python -m sim.scripts.dagger \
+    --repo-id local/safe-cube-dagger --root data/safe_cube_dagger \
+    --output-root outputs/safe_pi0_dagger \
+    --rounds 3 --episodes-per-round 40 \
+    --base-policy outputs/safe_pi0_bc/checkpoints/last/pretrained_model \
+    --safety-weight 1.0 --steps 8000 --batch-size 8
+```
+
+Use `--no-train` to only collect, or `--print-only` to see the constructed
+training commands without running them.
+
+### 4. Evaluate
+
+`sim.evaluate.evaluate()` reports success / red-contact / clearance / ceiling
+rates. Wrap a checkpoint with `sim.dagger.PolicyRollout` and pass its `.act` as
+the `policy` argument:
 
 ```python
-from sim import SafeCubeEnv, EnvConfig, ExpertConfig
-from sim.expert import ScriptedExpert
-from sim.recorder import EpisodeRecorder
-
-env = SafeCubeEnv(EnvConfig(...))
-expert = ScriptedExpert(env=env, cfg=ExpertConfig())
-
-for j in range(N_DAGGER_ROUNDS):
-    alpha = ALPHA0 ** j           # mixing weight: expert-heavy -> policy-heavy
-    for ep in range(EPISODES_PER_ROUND):
-        obs, info = env.reset(seed=...)
-        rec.begin_episode(task=...)
-        expert.reset()
-        while not done:
-            expert_action = expert.act(info)
-            policy_action = policy.act(obs)       # your π0 inference
-            use_expert = rng.random() < alpha
-            action = expert_action if use_expert else policy_action
-            # Always label with the expert action (DAgger).
-            rec.add(obs, expert_action, info)
-            obs, _, terminated, truncated, info = env.step(action)
-        rec.save_episode()
-    train_one_pass(policy, dataset, L_task + lam * L_safety)
+from sim import SafeCubeEnv, EnvConfig
+from sim.dagger import PolicyRollout
+from sim.evaluate import evaluate
+
+env = SafeCubeEnv(EnvConfig())
+roll = PolicyRollout(
+    checkpoint="outputs/safe_pi0_dagger/round_2/checkpoints/last/pretrained_model",
+    dataset_repo_id="local/safe-cube-dagger", dataset_root="data/safe_cube_dagger",
+)
+task = "pick up the blue cube ... place it on the goal patch"
+res = evaluate(env=env, policy=lambda obs: roll.act(obs, task), n_episodes=50)
+print(res.summary())
 ```
 
+### 5. Deploy
+
+Identical to a vanilla π0: `image → π0 → action`. No SDF, no cube positions, no
+filter — the safety preference lives in the weights. (On real hardware, keep the
+thin clearance watchdog from `project_summary.md` §7.4 as insurance.)
+
 ## Knobs worth tuning first
 
 | Knob | Where | Effect |