Auto-sync: 2026-06-27 15:04:36 (part 3)

scripts/eval_maniskill_policy_rollout.py

@@ -38,12 +38,22 @@ def main(argv: list[str] | None = None) -> int:
     )
     parser.add_argument(
         "--selection-mode",
-        choices=("policy", "field", "lattice", "retrieval_lattice"),
+        choices=(
+            "policy",
+            "field",
+            "field_optim",
+            "lattice",
+            "retrieval_lattice",
+            "retrieval_residual",
+        ),
         default="policy",
         help="'policy' executes the deterministic policy mean; 'field' scores model-generated "
-        "candidates with the learned interventional field; 'lattice' scores the current "
-        "state's CIL action lattice without reading rewards; 'retrieval_lattice' scores "
-        "the nearest train-state lattice for the current state.",
+        "candidates with the learned interventional field; 'field_optim' additionally "
+        "optimizes model-generated candidates with projected action-space gradient ascent; "
+        "'lattice' scores the current state's CIL action lattice without reading rewards; "
+        "'retrieval_lattice' scores the nearest train-state lattice for the current state; "
+        "'retrieval_residual' translates nearest train-state counterfactual residuals around "
+        "the current policy mean before field scoring.",
     )
     parser.add_argument(
         "--num-candidates",
@@ -63,6 +73,36 @@ def main(argv: list[str] | None = None) -> int:
         default=0,
         help="Base RNG seed for candidate perturbations.",
     )
+    parser.add_argument(
+        "--field-optim-steps",
+        type=int,
+        default=0,
+        help="Projected gradient-ascent steps when selection-mode=field_optim.",
+    )
+    parser.add_argument(
+        "--field-optim-step-size",
+        type=float,
+        default=0.05,
+        help="Action-space step size for selection-mode=field_optim.",
+    )
+    parser.add_argument(
+        "--field-optim-trust-radius",
+        type=float,
+        default=0.5,
+        help="L-infinity trust radius around the policy proposal for field_optim.",
+    )
+    parser.add_argument(
+        "--field-optim-l2-penalty",
+        type=float,
+        default=0.0,
+        help="Mean squared action penalty around the policy proposal for field_optim.",
+    )
+    parser.add_argument(
+        "--retrieval-neighbors",
+        type=int,
+        default=1,
+        help="Nearest train states to use for retrieval_lattice/retrieval_residual proposals.",
+    )
     parser.add_argument(
         "--lattice-exclude-types",
         default="",
@@ -87,6 +127,11 @@ def main(argv: list[str] | None = None) -> int:
         num_candidates=args.num_candidates,
         candidate_sigma=args.candidate_sigma,
         selection_seed=args.selection_seed,
+        field_optim_steps=args.field_optim_steps,
+        field_optim_step_size=args.field_optim_step_size,
+        field_optim_trust_radius=args.field_optim_trust_radius,
+        field_optim_l2_penalty=args.field_optim_l2_penalty,
+        retrieval_neighbors=args.retrieval_neighbors,
         lattice_exclude_types=lattice_exclude_types,
     )
     print(json.dumps({key: value for key, value in result.items() if key != "rows"}, indent=2))

scripts/slurm/eval_h16_field_sweep.sbatch

@@ -14,8 +14,9 @@
 set -euo pipefail
 
 # Field-guided h=16 rollout sweep.
-# Each job evaluates one seed/config pair by sampling model-generated action chunks,
-# scoring them with DoVLA's learned interventional field, and executing only the best.
+# Each job evaluates one seed/config pair by generating deploy-clean action chunks,
+# optionally optimizing them in action space with DoVLA's learned interventional field,
+# and executing only the best.
 
 PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
 SCRATCH_ROOT="/scratch/$USER/dovla"
@@ -43,6 +44,11 @@ DATASET="${DATASET:-$SCRATCH_ROOT/experiments/six_task_h16_collection}"
 SOURCE_RUN_ROOT="${SOURCE_RUN_ROOT:-$SCRATCH_ROOT/experiments/dovla_h16_rollout_runs}"
 SOURCE_OBJECTIVE="${SOURCE_OBJECTIVE:-}"
 CHECKPOINT_NAME="${CHECKPOINT_NAME:-best.pt}"
+SELECTION_MODE="${SELECTION_MODE:-field}"
+FIELD_OPTIM_STEPS="${FIELD_OPTIM_STEPS:-0}"
+FIELD_OPTIM_STEP_SIZE="${FIELD_OPTIM_STEP_SIZE:-0.05}"
+FIELD_OPTIM_TRUST_RADIUS="${FIELD_OPTIM_TRUST_RADIUS:-0.5}"
+FIELD_OPTIM_L2_PENALTY="${FIELD_OPTIM_L2_PENALTY:-0.0}"
 if [[ -z "${CHECKPOINT:-}" ]]; then
   if [[ -n "$SOURCE_OBJECTIVE" ]]; then
     CHECKPOINT="$SOURCE_RUN_ROOT/$SOURCE_OBJECTIVE/seed_$SEED/$CHECKPOINT_NAME"
@@ -75,6 +81,11 @@ echo "Seed: $SEED"
 echo "Candidates: $NUM_CANDIDATES"
 echo "Candidate sigma: $CANDIDATE_SIGMA"
 echo "Selection seed: $SELECTION_SEED"
+echo "Selection mode: $SELECTION_MODE"
+echo "Field optim steps: $FIELD_OPTIM_STEPS"
+echo "Field optim step size: $FIELD_OPTIM_STEP_SIZE"
+echo "Field optim trust radius: $FIELD_OPTIM_TRUST_RADIUS"
+echo "Field optim L2 penalty: $FIELD_OPTIM_L2_PENALTY"
 echo "Checkpoint: $CHECKPOINT"
 echo "Dataset: $DATASET"
 echo "Out: $OUT"
@@ -92,10 +103,14 @@ apptainer exec --nv --env "$ENVS" \
     --group-batch-size "$GROUP_BATCH_SIZE" \
     --sim-backend physx_cuda:0 \
     --render-backend cpu \
-    --selection-mode field \
+    --selection-mode "$SELECTION_MODE" \
     --num-candidates "$NUM_CANDIDATES" \
     --candidate-sigma "$CANDIDATE_SIGMA" \
-    --selection-seed "$SELECTION_SEED"
+    --selection-seed "$SELECTION_SEED" \
+    --field-optim-steps "$FIELD_OPTIM_STEPS" \
+    --field-optim-step-size "$FIELD_OPTIM_STEP_SIZE" \
+    --field-optim-trust-radius "$FIELD_OPTIM_TRUST_RADIUS" \
+    --field-optim-l2-penalty "$FIELD_OPTIM_L2_PENALTY"
 
 echo ""
 echo "Field-guided rollout complete"

scripts/slurm/eval_maniskill_policy_rollout.sbatch

@@ -44,6 +44,11 @@ SELECTION_MODE="${SELECTION_MODE:-policy}"
 NUM_CANDIDATES="${NUM_CANDIDATES:-1}"
 CANDIDATE_SIGMA="${CANDIDATE_SIGMA:-0.2}"
 SELECTION_SEED="${SELECTION_SEED:-0}"
+FIELD_OPTIM_STEPS="${FIELD_OPTIM_STEPS:-0}"
+FIELD_OPTIM_STEP_SIZE="${FIELD_OPTIM_STEP_SIZE:-0.05}"
+FIELD_OPTIM_TRUST_RADIUS="${FIELD_OPTIM_TRUST_RADIUS:-0.5}"
+FIELD_OPTIM_L2_PENALTY="${FIELD_OPTIM_L2_PENALTY:-0.0}"
+RETRIEVAL_NEIGHBORS="${RETRIEVAL_NEIGHBORS:-1}"
 LATTICE_EXCLUDE_TYPES="${LATTICE_EXCLUDE_TYPES:-}"
 if [[ -n "${LATTICE_EXCLUDE_TYPES_COLON:-}" ]]; then
   LATTICE_EXCLUDE_TYPES="${LATTICE_EXCLUDE_TYPES_COLON//:/,}"
@@ -87,5 +92,10 @@ apptainer exec --nv \
     --num-candidates "$NUM_CANDIDATES" \
     --candidate-sigma "$CANDIDATE_SIGMA" \
     --selection-seed "$SELECTION_SEED" \
+    --field-optim-steps "$FIELD_OPTIM_STEPS" \
+    --field-optim-step-size "$FIELD_OPTIM_STEP_SIZE" \
+    --field-optim-trust-radius "$FIELD_OPTIM_TRUST_RADIUS" \
+    --field-optim-l2-penalty "$FIELD_OPTIM_L2_PENALTY" \
+    --retrieval-neighbors "$RETRIEVAL_NEIGHBORS" \
     --lattice-exclude-types "$LATTICE_EXCLUDE_TYPES" \
     "${EXTRA_ARGS[@]}"

scripts/slurm/summarize_h16_field_sweep.sbatch

@@ -12,13 +12,7 @@ set -euo pipefail
 PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
 RUN_ROOT="${RUN_ROOT:-/scratch/$USER/dovla/experiments/dovla_h16_field_sweep}"
 SUMMARY_TAG="${SUMMARY_TAG:-field_sweep}"
-if [[ -z "${PYTHON:-}" ]]; then
-  if [[ -x "$PROJECT_DIR/.venv/bin/python" ]]; then
-    PYTHON="$PROJECT_DIR/.venv/bin/python"
-  else
-    PYTHON="python3"
-  fi
-fi
+PYTHON="${PYTHON:-python3}"
 export RUN_ROOT
 export SUMMARY_TAG
 
@@ -52,6 +46,11 @@ for result_path in sorted(run_root.glob("k*_sigma*/seed_*/online_rollout.json"))
             "selection_mode": data.get("selection_mode"),
             "num_candidates": data.get("num_candidates"),
             "candidate_sigma": data.get("candidate_sigma"),
+            "field_optim_steps": data.get("field_optim_steps", 0),
+            "field_optim_step_size": data.get("field_optim_step_size", 0.0),
+            "field_optim_trust_radius": data.get("field_optim_trust_radius", 0.0),
+            "field_optim_l2_penalty": data.get("field_optim_l2_penalty", 0.0),
+            "retrieval_neighbors": data.get("retrieval_neighbors", 0),
             "num_groups": data.get("num_groups"),
             "policy_rollout_success_rate": data.get("policy_rollout_success_rate", 0.0),
             "policy_rollout_progress": data.get("policy_rollout_progress", 0.0),

scripts/slurm/summarize_h16_lattice_rollout.sbatch

@@ -12,13 +12,7 @@ set -euo pipefail
 PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
 RUN_ROOT="${RUN_ROOT:-/scratch/$USER/dovla/experiments/dovla_h16_rollout_runs}"
 OUT_NAME="${OUT_NAME:-lattice_rollout.json}"
-if [[ -z "${PYTHON:-}" ]]; then
-  if [[ -x "$PROJECT_DIR/.venv/bin/python" ]]; then
-    PYTHON="$PROJECT_DIR/.venv/bin/python"
-  else
-    PYTHON="python3"
-  fi
-fi
+PYTHON="${PYTHON:-python3}"
 export RUN_ROOT
 export OUT_NAME
 

scripts/slurm/summarize_h16_policy_ckpt.sbatch

@@ -14,13 +14,7 @@ RUN_ROOT="${RUN_ROOT:-/scratch/$USER/dovla/experiments/dovla_h16_policy_ckpt_run
 OBJECTIVE="${OBJECTIVE:-base}"
 OUT_NAME="${OUT_NAME:-policy_rollout.json}"
 SUMMARY_TAG="${SUMMARY_TAG:-}"
-if [[ -z "${PYTHON:-}" ]]; then
-  if [[ -x "$PROJECT_DIR/.venv/bin/python" ]]; then
-    PYTHON="$PROJECT_DIR/.venv/bin/python"
-  else
-    PYTHON="python3"
-  fi
-fi
+PYTHON="${PYTHON:-python3}"
 export RUN_ROOT
 export OBJECTIVE
 export OUT_NAME
@@ -61,6 +55,11 @@ for result_path in sorted(base_dir.glob(f"seed_*/{out_name}")):
             "selection_mode": data.get("selection_mode"),
             "num_candidates": data.get("num_candidates"),
             "candidate_sigma": data.get("candidate_sigma"),
+            "field_optim_steps": data.get("field_optim_steps", 0),
+            "field_optim_step_size": data.get("field_optim_step_size", 0.0),
+            "field_optim_trust_radius": data.get("field_optim_trust_radius", 0.0),
+            "field_optim_l2_penalty": data.get("field_optim_l2_penalty", 0.0),
+            "retrieval_neighbors": data.get("retrieval_neighbors", 0),
             "policy_rollout_success_rate": data.get("policy_rollout_success_rate", 0.0),
             "policy_rollout_progress": data.get("policy_rollout_progress", 0.0),
             "oracle_success_rate": data.get("oracle_success_rate", 0.0),
@@ -121,17 +120,20 @@ lines = [
     f"Mean progress: {summary['mean_progress']:.2%}",
     f"Mean action MSE to best: {summary['mean_action_mse_to_best']:.3f}",
     "",
-    "| seed | mode | k | sigma | success | progress | oracle | action MSE |",
-    "|---:|---|---:|---:|---:|---:|---:|---:|",
+    "| seed | mode | k | retrieval K | sigma | opt steps | trust | success | progress | oracle | action MSE |",
+    "|---:|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|",
 ]
 for row in rows:
     lines.append(
-        "| {seed} | {mode} | {k} | {sigma:.2f} | {success:.2%} | "
-        "{progress:.2%} | {oracle:.2%} | {mse:.3f} |".format(
+        "| {seed} | {mode} | {k} | {retrieval} | {sigma:.2f} | {steps} | {trust:.2f} | "
+        "{success:.2%} | {progress:.2%} | {oracle:.2%} | {mse:.3f} |".format(
             seed=row["seed"],
             mode=row.get("selection_mode") or "policy",
             k=row.get("num_candidates") or 1,
+            retrieval=row.get("retrieval_neighbors") or 0,
             sigma=row.get("candidate_sigma") or 0.0,
+            steps=row.get("field_optim_steps") or 0,
+            trust=row.get("field_optim_trust_radius") or 0.0,
             success=row["policy_rollout_success_rate"],
             progress=row["policy_rollout_progress"],
             oracle=row["oracle_success_rate"],

tests/test_maniskill_policy_rollout.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import pickle
 from pathlib import Path
+from types import SimpleNamespace
 
 import numpy as np
 
@@ -14,7 +15,9 @@ from dovla_cil.data.schema import (
     StructuredEffect,
 )
 from dovla_cil.eval.maniskill_policy_rollout import (
+    _RolloutCase,
     _adapt_action_dim,
+    _attach_retrieved_residual_candidates,
     _load_state_archive,
     _numeric_action_values,
     _select_action_chunk,
@@ -220,6 +223,60 @@ def test_field_mode_scores_clamped_candidates() -> None:
     assert index.tolist()[0] != 0
 
 
+def test_field_optim_mode_improves_policy_mean() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 3)
+    offset = torch.full_like(mean, 0.4)
+    model = _StubModel(torch, mean, best_offset=offset)
+    actions, index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="field_optim",
+        num_candidates=1,
+        candidate_sigma=0.0,
+        selection_seed=7,
+        field_optim_steps=8,
+        field_optim_step_size=0.1,
+        field_optim_trust_radius=0.5,
+        field_optim_l2_penalty=0.0,
+    )
+
+    assert float(((actions - (mean + offset)) ** 2).sum()) < float(
+        ((mean - (mean + offset)) ** 2).sum()
+    )
+    assert index.tolist() == [0]
+
+
+def test_field_optim_mode_respects_trust_region_and_bounds() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 3)
+    offset = torch.full_like(mean, 10.0)
+    model = _StubModel(torch, mean, best_offset=offset)
+    actions, _index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="field_optim",
+        num_candidates=4,
+        candidate_sigma=1.0,
+        selection_seed=7,
+        field_optim_steps=8,
+        field_optim_step_size=0.2,
+        field_optim_trust_radius=0.25,
+        field_optim_l2_penalty=0.0,
+        action_low=torch.full_like(mean, -0.5),
+        action_high=torch.full_like(mean, 0.5),
+    )
+
+    assert float(actions.max()) <= 0.25
+    assert float(actions.min()) >= -0.25
+
+
 def test_lattice_mode_selects_best_scored_candidate() -> None:
     import torch
 
@@ -288,3 +345,96 @@ def test_retrieval_lattice_mode_uses_candidate_tensor() -> None:
 
     assert torch.allclose(actions, mean + offset)
     assert index.tolist() == [1]
+
+
+def test_retrieval_residual_mode_translates_residuals_around_policy_mean() -> None:
+    import torch
+
+    mean = torch.full((1, 1, 3), 0.1)
+    offset = torch.full_like(mean, 0.4)
+    model = _StubModel(torch, mean, best_offset=offset)
+    residuals = torch.stack([torch.zeros_like(mean), offset], dim=1)
+    actions, index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="retrieval_residual",
+        num_candidates=1,
+        candidate_sigma=0.0,
+        selection_seed=0,
+        action_candidates=residuals,
+    )
+
+    assert torch.allclose(actions, mean + offset)
+    assert index.tolist() == [1]
+
+
+def test_retrieval_residual_candidates_use_knn_train_residuals() -> None:
+    def record(group_id: str, candidate_type: str, action_value: float, feature: float):
+        return SimpleNamespace(
+            group_id=group_id,
+            task_id="PickCube-v1",
+            candidate_type=candidate_type,
+            record_id=f"{group_id}-{candidate_type}-{action_value}",
+            observation_inline={"features": [feature, 0.0]},
+            action_chunk=ActionChunk(
+                representation="continuous",
+                horizon=1,
+                values=[[action_value, 0.0]],
+            ),
+        )
+
+    dataset = SimpleNamespace(
+        group_ids=["train_a", "train_b", "heldout"],
+        get_group=lambda group_id: {
+            "train_a": [
+                record("train_a", "expert", 1.0, 0.0),
+                record("train_a", "near_miss", 1.2, 0.0),
+            ],
+            "train_b": [
+                record("train_b", "expert", -1.0, 0.4),
+                record("train_b", "wrong_direction", -1.3, 0.4),
+            ],
+            "heldout": [
+                record("heldout", "expert", 9.0, 0.1),
+                record("heldout", "near_miss", 9.9, 0.1),
+            ],
+        }[group_id],
+    )
+    case = _RolloutCase(
+        group_id="heldout",
+        task_id="PickCube-v1",
+        source_dataset=Path("."),
+        state={},
+        observation={"features": [0.1, 0.0]},
+        instruction="pick",
+        oracle_score=1.0,
+        oracle_success=True,
+        expert_score=1.0,
+        expert_success=True,
+        best_action_values=[[9.9, 0.0]],
+        candidate_action_values=[],
+        candidate_types=[],
+    )
+
+    [attached] = _attach_retrieved_residual_candidates(
+        dataset,
+        [case],
+        heldout_group_ids=["heldout"],
+        obs_dim=2,
+        observation_mode="state",
+        retrieval_neighbors=2,
+    )
+
+    assert attached.candidate_source_group_id == "train_a;train_b"
+    assert attached.candidate_types == [
+        "policy_residual",
+        "residual_near_miss",
+        "policy_residual",
+        "residual_wrong_direction",
+    ]
+    assert np.allclose(
+        np.asarray(attached.candidate_action_values, dtype=np.float32),
+        np.asarray([[[0.0, 0.0]], [[0.2, 0.0]], [[0.0, 0.0]], [[-0.3, 0.0]]]),
+    )