Add files using upload-large-folder tool

code/reveal_vla_bimanual/eval/metrics.py

@@ -22,6 +22,13 @@ class PlannerDiagnostics:
     regret: float
     risk_calibration_mse: float
     role_collapse_rate: float
+    proposal_diversity: float | None = None
+    planner_score_utility_spearman: float | None = None
+    left_right_equivariance_error: float | None = None
+    belief_calibration_brier: float | None = None
+    reocclusion_calibration_brier: float | None = None
+    support_stability_mae: float | None = None
+    clearance_auc: float | None = None
 
 
 def mean_success(per_task_success: dict[str, float]) -> float:
@@ -87,6 +94,84 @@ def risk_calibration_mse(predicted_risk: np.ndarray, realized_risk: np.ndarray)
     return float(np.mean((predicted_risk - realized_risk) ** 2))
 
 
+def proposal_diversity(proposal_chunks: np.ndarray) -> float:
+    proposal_chunks = np.asarray(proposal_chunks, dtype=np.float32)
+    if proposal_chunks.ndim != 4 or proposal_chunks.shape[1] <= 1:
+        return 0.0
+    flat = proposal_chunks.reshape(proposal_chunks.shape[0], proposal_chunks.shape[1], -1)
+    diffs = flat[:, :, None, :] - flat[:, None, :, :]
+    distances = np.abs(diffs).mean(axis=-1)
+    mask = ~np.eye(distances.shape[1], dtype=bool)
+    if not mask.any():
+        return 0.0
+    off_diagonal = distances[:, mask]
+    return float(off_diagonal.mean())
+
+
+def planner_score_utility_spearman(pred_scores: np.ndarray, oracle_utility: np.ndarray) -> float:
+    pred_scores = np.asarray(pred_scores, dtype=np.float32)
+    oracle_utility = np.asarray(oracle_utility, dtype=np.float32)
+    if pred_scores.size == 0:
+        return 0.0
+    pred_rank = pred_scores.argsort(axis=-1).argsort(axis=-1).astype(np.float32)
+    oracle_rank = oracle_utility.argsort(axis=-1).argsort(axis=-1).astype(np.float32)
+    pred_rank = pred_rank - pred_rank.mean(axis=-1, keepdims=True)
+    oracle_rank = oracle_rank - oracle_rank.mean(axis=-1, keepdims=True)
+    denom = np.sqrt((pred_rank**2).sum(axis=-1) * (oracle_rank**2).sum(axis=-1))
+    valid = denom > 1e-6
+    if not np.any(valid):
+        return 0.0
+    corr = np.zeros_like(denom)
+    corr[valid] = (pred_rank[valid] * oracle_rank[valid]).sum(axis=-1) / denom[valid]
+    return float(corr.mean())
+
+
+def left_right_equivariance_error(pred: np.ndarray, swapped_target: np.ndarray) -> float:
+    pred = np.asarray(pred, dtype=np.float32)
+    swapped_target = np.asarray(swapped_target, dtype=np.float32)
+    if pred.size == 0 or swapped_target.size == 0:
+        return 0.0
+    return float(np.abs(pred - swapped_target).mean())
+
+
+def belief_calibration_brier(predicted_belief: np.ndarray, target_belief: np.ndarray) -> float:
+    predicted_belief = np.asarray(predicted_belief, dtype=np.float32)
+    target_belief = np.asarray(target_belief, dtype=np.float32)
+    if predicted_belief.size == 0:
+        return 0.0
+    return float(np.mean((predicted_belief - target_belief) ** 2))
+
+
+def reocclusion_calibration_brier(predicted_reocclusion: np.ndarray, target_reocclusion: np.ndarray) -> float:
+    predicted_reocclusion = np.asarray(predicted_reocclusion, dtype=np.float32)
+    target_reocclusion = np.asarray(target_reocclusion, dtype=np.float32)
+    if predicted_reocclusion.size == 0:
+        return 0.0
+    return float(np.mean((predicted_reocclusion - target_reocclusion) ** 2))
+
+
+def support_stability_mae(predicted: np.ndarray, target: np.ndarray) -> float:
+    predicted = np.asarray(predicted, dtype=np.float32)
+    target = np.asarray(target, dtype=np.float32)
+    if predicted.size == 0:
+        return 0.0
+    return float(np.abs(predicted - target).mean())
+
+
+def clearance_auc(predicted: np.ndarray, target: np.ndarray) -> float:
+    predicted = np.asarray(predicted, dtype=np.float32).reshape(-1)
+    target = np.asarray(target, dtype=np.float32).reshape(-1)
+    positives = target > 0.5
+    negatives = ~positives
+    if positives.sum() == 0 or negatives.sum() == 0:
+        return 0.0
+    order = np.argsort(predicted)
+    ranks = np.empty_like(order, dtype=np.float32)
+    ranks[order] = np.arange(order.shape[0], dtype=np.float32)
+    pos_ranks = ranks[positives]
+    return float((pos_ranks.sum() - positives.sum() * (positives.sum() - 1) / 2.0) / (positives.sum() * negatives.sum()))
+
+
 def role_collapse_rate(
     action_chunks: np.ndarray,
     arm_role_logits: np.ndarray | None = None,

code/reveal_vla_bimanual/eval/run_peract2_launch_smoke.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+from sim_rlbench.task_splits import PERACT2_BIMANUAL_TASKS
+
+
+def _parse_json_payload(stdout: str) -> dict[str, Any]:
+    start = stdout.find("{")
+    end = stdout.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        raise ValueError("No JSON object found in subprocess stdout.")
+    return json.loads(stdout[start : end + 1])
+
+
+def _run_task(project_root: Path, output_dir: Path, task_name: str, *, resolution: int, headless: bool) -> dict[str, Any]:
+    task_dir = output_dir / task_name
+    task_dir.mkdir(parents=True, exist_ok=True)
+    command = [
+        sys.executable,
+        "-m",
+        "sim_rlbench.launch_smoke",
+        "--task",
+        task_name,
+        "--resolution",
+        str(resolution),
+    ]
+    if headless:
+        command.append("--headless")
+    completed = subprocess.run(
+        command,
+        cwd=project_root,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    (task_dir / "command.txt").write_text(" ".join(command) + "\n", encoding="utf-8")
+    (task_dir / "stdout.txt").write_text(completed.stdout, encoding="utf-8")
+    (task_dir / "stderr.txt").write_text(completed.stderr, encoding="utf-8")
+
+    payload: dict[str, Any] = {
+        "subprocess_returncode": int(completed.returncode),
+        "launch_ok": completed.returncode == 0,
+    }
+    try:
+        payload.update(_parse_json_payload(completed.stdout))
+    except Exception as exc:
+        payload["launch_ok"] = False
+        payload["error"] = f"json_parse_failed: {exc}"
+    if completed.returncode != 0 and "error" not in payload:
+        payload["error"] = f"subprocess_exit_{completed.returncode}"
+    return payload
+
+
+def _write_markdown(path: Path, payload: dict[str, Any]) -> None:
+    lines = [
+        "# PerAct2 13-Task Launch Smoke",
+        "",
+        f"- Resolution: `{payload['resolution']}`",
+        f"- Headless: `{payload['headless']}`",
+        f"- Task count: `{payload['task_count']}`",
+        f"- Launch successes: `{payload['launch_successes']}`",
+        f"- Finite-action tasks: `{payload['finite_action_tasks']}`",
+        f"- Error tasks: `{payload['error_tasks']}`",
+        "",
+        "## Per-task",
+        "",
+    ]
+    for task_name, task_payload in payload["tasks"].items():
+        if "error" in task_payload:
+            lines.append(
+                f"- `{task_name}`: launch_ok={task_payload.get('launch_ok')}, "
+                f"action_finite={task_payload.get('action_finite')}, "
+                f"error={task_payload['error']}, "
+                f"subprocess_returncode={task_payload['subprocess_returncode']}"
+            )
+        else:
+            lines.append(
+                f"- `{task_name}`: launch_ok={task_payload.get('launch_ok')}, "
+                f"action_finite={task_payload.get('action_finite')}, "
+                f"task_class={task_payload.get('task')}, "
+                f"reward={task_payload.get('reward')}"
+            )
+    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--tasks", nargs="*", default=list(PERACT2_BIMANUAL_TASKS))
+    parser.add_argument("--resolution", type=int, default=224)
+    parser.add_argument("--headless", action="store_true", default=True)
+    args = parser.parse_args()
+
+    project_root = Path(__file__).resolve().parents[1]
+    output_dir = Path(args.output_dir).resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    results: dict[str, Any] = {
+        "resolution": int(args.resolution),
+        "headless": bool(args.headless),
+        "tasks": {},
+    }
+    for task_name in tuple(args.tasks):
+        print(f"[peract2-launch-smoke] task={task_name}", flush=True)
+        results["tasks"][task_name] = _run_task(
+            project_root,
+            output_dir,
+            task_name,
+            resolution=args.resolution,
+            headless=args.headless,
+        )
+
+    task_payloads = list(results["tasks"].values())
+    results["task_count"] = len(task_payloads)
+    results["launch_successes"] = int(sum(1 for payload in task_payloads if payload.get("launch_ok")))
+    results["finite_action_tasks"] = int(sum(1 for payload in task_payloads if payload.get("action_finite")))
+    results["error_tasks"] = sorted(task_name for task_name, payload in results["tasks"].items() if "error" in payload)
+
+    (output_dir / "launch_smoke_summary.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
+    _write_markdown(output_dir / "launch_smoke_summary.md", results)
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()

code/reveal_vla_bimanual/eval/run_proxy_diagnostics.py

@@ -8,9 +8,22 @@ from typing import Any
 import numpy as np
 import torch
 from torch import Tensor
+import torch.nn.functional as F
 from torch.utils.data import DataLoader
 
-from eval.metrics import planner_regret, planner_top1_accuracy, risk_calibration_mse, role_collapse_rate
+from eval.metrics import (
+    belief_calibration_brier,
+    clearance_auc,
+    left_right_equivariance_error,
+    planner_regret,
+    planner_score_utility_spearman,
+    planner_top1_accuracy,
+    proposal_diversity,
+    reocclusion_calibration_brier,
+    risk_calibration_mse,
+    role_collapse_rate,
+    support_stability_mae,
+)
 from eval.run_reveal_benchmark import load_model
 from sim_reveal.dataset import dataset_from_bundle, load_teacher_dataset
 
@@ -52,41 +65,121 @@ def main() -> None:
     risk_batches: list[np.ndarray] = []
     realized_risk_batches: list[np.ndarray] = []
     collapse_batches: list[float] = []
+    proposal_batches: list[np.ndarray] = []
+    equivariance_batches: list[float] = []
+    belief_pred_batches: list[np.ndarray] = []
+    belief_target_batches: list[np.ndarray] = []
+    reocclusion_pred_batches: list[np.ndarray] = []
+    reocclusion_target_batches: list[np.ndarray] = []
+    support_pred_batches: list[np.ndarray] = []
+    support_target_batches: list[np.ndarray] = []
+    clearance_pred_batches: list[np.ndarray] = []
+    clearance_target_batches: list[np.ndarray] = []
+    memory_write_batches: list[np.ndarray] = []
+    memory_saturation_batches: list[np.ndarray] = []
 
     with torch.no_grad():
         for batch in loader:
             moved = _move_batch_to_device(batch, device)
-            outputs = model(
-                images=moved["images"],
-                proprio=moved["proprio"],
-                texts=moved["texts"],
-                history_images=moved.get("history_images"),
-                history_proprio=moved.get("history_proprio"),
-                history_actions=moved.get("history_actions"),
-                plan=True,
-                candidate_chunks_override=moved["candidate_action_chunks"],
-            )
+            forward_kwargs = {
+                "images": moved["images"],
+                "proprio": moved["proprio"],
+                "texts": moved["texts"],
+                "history_images": moved.get("history_images"),
+                "history_proprio": moved.get("history_proprio"),
+                "history_actions": moved.get("history_actions"),
+                "plan": True,
+                "candidate_chunks_override": moved["candidate_action_chunks"],
+            }
+            if hasattr(model, "elastic_state_head"):
+                forward_kwargs.update(
+                    {
+                        "depths": moved.get("depths"),
+                        "depth_valid": moved.get("depth_valid"),
+                        "camera_intrinsics": moved.get("camera_intrinsics"),
+                        "camera_extrinsics": moved.get("camera_extrinsics"),
+                        "history_depths": moved.get("history_depths"),
+                        "history_depth_valid": moved.get("history_depth_valid"),
+                        "use_depth": moved.get("depths") is not None,
+                        "use_world_model": True,
+                        "use_planner": True,
+                        "use_role_tokens": True,
+                        "compute_equivariance_probe": True,
+                    }
+                )
+            outputs = model(**forward_kwargs)
             if "planner_scores" not in outputs:
                 raise RuntimeError("Planner outputs were not produced for proxy diagnostics.")
-            score_batches.append(outputs["planner_scores"].detach().cpu().numpy())
-            utility_batches.append(moved["candidate_utility"].detach().cpu().numpy())
-            best_index_batches.append(outputs["best_candidate_indices"].detach().cpu().numpy())
-            risk_batches.append(outputs["planner_risk_values"].detach().cpu().numpy())
-            realized_risk_batches.append(
-                torch.clamp(
-                    moved["candidate_final_disturbance_cost"] + moved["candidate_reocclusion_rate"],
-                    0.0,
-                    1.0,
-                )
-                .detach()
-                .cpu()
-                .numpy()
+            planner_scores = outputs["planner_scores"]
+            candidate_utility = moved["candidate_utility"]
+            predicted_risk = outputs["planner_risk_values"]
+            realized_risk = torch.clamp(
+                moved["candidate_final_disturbance_cost"] + moved["candidate_reocclusion_rate"],
+                0.0,
+                1.0,
             )
+            shortlist_indices = outputs.get("planner_topk_indices")
+            if shortlist_indices is not None:
+                candidate_utility = candidate_utility.gather(1, shortlist_indices)
+                predicted_risk = predicted_risk
+                realized_risk = realized_risk.gather(1, shortlist_indices)
+            score_batches.append(planner_scores.detach().cpu().numpy())
+            utility_batches.append(candidate_utility.detach().cpu().numpy())
+            best_index_batches.append(outputs["best_candidate_indices"].detach().cpu().numpy())
+            risk_batches.append(predicted_risk.detach().cpu().numpy())
+            realized_risk_batches.append(realized_risk.detach().cpu().numpy())
             selected_chunk = outputs["planned_chunk"].detach().cpu().numpy()[:, None]
+            state = outputs.get("interaction_state") or outputs.get("reveal_state")
             role_logits = None
-            if outputs.get("interaction_state") is not None:
-                role_logits = outputs["interaction_state"]["arm_role_logits"].detach().cpu().numpy()[:, None]
+            if state is not None:
+                role_logits = state["arm_role_logits"].detach().cpu().numpy()[:, None]
             collapse_batches.append(role_collapse_rate(selected_chunk, role_logits))
+            if outputs.get("proposal_candidates") is not None:
+                proposal_batches.append(outputs["proposal_candidates"].detach().cpu().numpy())
+            if outputs.get("equivariance_probe_action_mean") is not None:
+                equivariance_batches.append(
+                    left_right_equivariance_error(
+                        outputs["equivariance_probe_action_mean"].detach().cpu().numpy(),
+                        outputs["equivariance_target_action_mean"].detach().cpu().numpy(),
+                    )
+                )
+            if state is not None:
+                if "belief_map" in state and "belief_map" in moved:
+                    belief_pred_batches.append(torch.sigmoid(state["belief_map"]).detach().cpu().numpy())
+                    belief_target_batches.append(moved["belief_map"].detach().cpu().numpy())
+                if "reocclusion_field" in state and "reocclusion_target" in moved:
+                    reocclusion_pred_batches.append(torch.sigmoid(state["reocclusion_field"]).mean(dim=(-1, -2)).detach().cpu().numpy())
+                    reocclusion_target_batches.append(moved["reocclusion_target"].detach().cpu().numpy())
+                if "support_stability_field" in state and "support_stability" in moved:
+                    support_pred_batches.append(torch.sigmoid(state["support_stability_field"]).mean(dim=(-1, -2)).detach().cpu().numpy())
+                    support_target_batches.append(moved["support_stability"].detach().cpu().numpy())
+                if "clearance_field" in state and "clearance_map" in moved:
+                    clearance_pred = torch.sigmoid(state["clearance_field"])
+                    clearance_target = moved["clearance_map"]
+                    if clearance_pred.shape[-2:] != clearance_target.shape[-2:]:
+                        clearance_pred = F.interpolate(
+                            clearance_pred,
+                            size=clearance_target.shape[-2:],
+                            mode="bilinear",
+                            align_corners=False,
+                        )
+                    if clearance_pred.shape[1] != clearance_target.shape[1]:
+                        if clearance_pred.shape[1] == 1:
+                            clearance_pred = clearance_pred.expand(-1, clearance_target.shape[1], -1, -1)
+                        elif clearance_target.shape[1] == 1:
+                            clearance_target = clearance_target.expand_as(clearance_pred)
+                        else:
+                            min_channels = min(clearance_pred.shape[1], clearance_target.shape[1])
+                            clearance_pred = clearance_pred[:, :min_channels]
+                            clearance_target = clearance_target[:, :min_channels]
+                    clearance_pred_batches.append(clearance_pred.detach().cpu().numpy())
+                    clearance_target_batches.append(clearance_target.detach().cpu().numpy())
+            if outputs.get("memory_output") is not None:
+                memory_output = outputs["memory_output"]
+                if "memory_write_rate" in memory_output:
+                    memory_write_batches.append(memory_output["memory_write_rate"].detach().cpu().numpy())
+                if "memory_saturation" in memory_output:
+                    memory_saturation_batches.append(memory_output["memory_saturation"].detach().cpu().numpy())
 
     scores = np.concatenate(score_batches, axis=0) if score_batches else np.zeros((0, 0), dtype=np.float32)
     utility = np.concatenate(utility_batches, axis=0) if utility_batches else np.zeros((0, 0), dtype=np.float32)
@@ -101,8 +194,37 @@ def main() -> None:
     diagnostics = {
         "planner_top1_accuracy": planner_top1_accuracy(scores, utility),
         "planner_regret": planner_regret(selected_indices, utility),
+        "planner_score_utility_spearman": planner_score_utility_spearman(scores, utility),
         "risk_calibration_mse": risk_calibration_mse(predicted_risk, realized_risk),
         "role_collapse_rate": float(np.mean(collapse_batches)) if collapse_batches else 0.0,
+        "proposal_diversity": proposal_diversity(np.concatenate(proposal_batches, axis=0)) if proposal_batches else 0.0,
+        "left_right_equivariance_error": float(np.mean(equivariance_batches)) if equivariance_batches else 0.0,
+        "belief_calibration_brier": belief_calibration_brier(
+            np.concatenate(belief_pred_batches, axis=0),
+            np.concatenate(belief_target_batches, axis=0),
+        )
+        if belief_pred_batches
+        else 0.0,
+        "reocclusion_calibration_brier": reocclusion_calibration_brier(
+            np.concatenate(reocclusion_pred_batches, axis=0),
+            np.concatenate(reocclusion_target_batches, axis=0),
+        )
+        if reocclusion_pred_batches
+        else 0.0,
+        "support_stability_mae": support_stability_mae(
+            np.concatenate(support_pred_batches, axis=0),
+            np.concatenate(support_target_batches, axis=0),
+        )
+        if support_pred_batches
+        else 0.0,
+        "clearance_auc": clearance_auc(
+            np.concatenate(clearance_pred_batches, axis=0),
+            np.concatenate(clearance_target_batches, axis=0),
+        )
+        if clearance_pred_batches
+        else 0.0,
+        "memory_write_rate": float(np.mean(np.concatenate(memory_write_batches, axis=0))) if memory_write_batches else 0.0,
+        "memory_saturation": float(np.mean(np.concatenate(memory_saturation_batches, axis=0))) if memory_saturation_batches else 0.0,
         "num_samples": int(scores.shape[0]),
     }
 

code/reveal_vla_bimanual/eval/run_reveal_benchmark.py

@@ -73,12 +73,18 @@ def _prepare_batch(
     observation: dict[str, Any],
     device: torch.device,
     history_images: list[np.ndarray] | None = None,
+    history_depths: list[np.ndarray] | None = None,
+    history_depth_valid: list[np.ndarray] | None = None,
     history_proprio: list[np.ndarray] | None = None,
     history_actions: list[np.ndarray] | None = None,
 ) -> dict[str, Any]:
     images = torch.from_numpy(observation["images"]).permute(0, 3, 1, 2).unsqueeze(0).float() / 255.0
+    depths = torch.from_numpy(observation.get("depths", np.zeros((3, 1, images.shape[-2], images.shape[-1]), dtype=np.float32))).unsqueeze(0).float()
+    depth_valid = torch.from_numpy(observation.get("depth_valid", np.zeros((3, 1, images.shape[-2], images.shape[-1]), dtype=np.float32))).unsqueeze(0).float()
     proprio = torch.from_numpy(observation["proprio"]).unsqueeze(0).float()
     history_images = history_images or []
+    history_depths = history_depths or []
+    history_depth_valid = history_depth_valid or []
     history_proprio = history_proprio or []
     history_actions = history_actions or []
     if history_images:
@@ -90,6 +96,12 @@ def _prepare_batch(
             (1, 0, images.shape[1], images.shape[2], images.shape[3], images.shape[4]),
             dtype=torch.float32,
         )
+    if history_depths:
+        history_depths_tensor = torch.from_numpy(np.stack(history_depths, axis=0)).unsqueeze(0).float()
+        history_depth_valid_tensor = torch.from_numpy(np.stack(history_depth_valid, axis=0)).unsqueeze(0).float()
+    else:
+        history_depths_tensor = torch.zeros((1, 0, depths.shape[1], depths.shape[2], depths.shape[3], depths.shape[4]), dtype=torch.float32)
+        history_depth_valid_tensor = torch.zeros_like(history_depths_tensor)
     if history_proprio:
         history_proprio_tensor = torch.from_numpy(np.stack(history_proprio, axis=0)).unsqueeze(0).float()
     else:
@@ -100,7 +112,13 @@ def _prepare_batch(
         history_actions_tensor = torch.zeros((1, 0, 14), dtype=torch.float32)
     return {
         "images": images.to(device),
+        "depths": depths.to(device),
+        "depth_valid": depth_valid.to(device),
+        "camera_intrinsics": torch.from_numpy(observation.get("camera_intrinsics", np.zeros((3, 3, 3), dtype=np.float32))).unsqueeze(0).to(device),
+        "camera_extrinsics": torch.from_numpy(observation.get("camera_extrinsics", np.zeros((3, 4, 4), dtype=np.float32))).unsqueeze(0).to(device),
         "history_images": history_images_tensor.to(device),
+        "history_depths": history_depths_tensor.to(device),
+        "history_depth_valid": history_depth_valid_tensor.to(device),
         "history_proprio": history_proprio_tensor.to(device),
         "history_actions": history_actions_tensor.to(device),
         "proprio": proprio.to(device),
@@ -147,6 +165,27 @@ def select_chunk(
         if "planned_chunk" in outputs and ablation not in {"no_world_model", "no_interaction_head"}:
             return outputs["planned_chunk"], outputs
         return outputs["action_mean"], outputs
+    if hasattr(model, "elastic_state_head"):
+        outputs = model(
+            **forward_kwargs,
+            depths=batch.get("depths"),
+            depth_valid=batch.get("depth_valid"),
+            camera_intrinsics=batch.get("camera_intrinsics"),
+            camera_extrinsics=batch.get("camera_extrinsics"),
+            history_depths=batch.get("history_depths"),
+            history_depth_valid=batch.get("history_depth_valid"),
+            plan=True,
+            use_world_model=(ablation not in {"no_world_model", "no_planner"}),
+            use_planner=(ablation != "no_planner"),
+            use_depth=(ablation != "no_depth"),
+            use_role_tokens=(ablation not in {"no_role_tokens", "no_role_symmetry"}),
+            history_steps_override=(2 if ablation == "short_history" else None),
+        )
+        if "planned_chunk" in outputs and ablation != "no_planner":
+            return outputs["planned_chunk"], outputs
+        if "candidate_chunks" in outputs:
+            return outputs["candidate_chunks"][:, 0], outputs
+        return outputs["action_mean"], outputs
     if hasattr(model, "reveal_head"):
         if ablation == "no_world_model":
             outputs = model(**forward_kwargs, plan=False)
@@ -195,6 +234,8 @@ def evaluate_model(
             episode_corridor = [float(privileged_state["corridor_feasible"][privileged_state["support_mode"]].any())]
             episode_disturbance = [float(privileged_state["disturbance_cost"])]
             history_images: list[np.ndarray] = []
+            history_depths: list[np.ndarray] = []
+            history_depth_valid: list[np.ndarray] = []
             history_proprio: list[np.ndarray] = []
             history_actions: list[np.ndarray] = []
             done = False
@@ -203,6 +244,8 @@ def evaluate_model(
                     observation,
                     device=device,
                     history_images=history_images,
+                    history_depths=history_depths,
+                    history_depth_valid=history_depth_valid,
                     history_proprio=history_proprio,
                     history_actions=history_actions,
                 )
@@ -224,9 +267,14 @@ def evaluate_model(
                     if history_steps > 0:
                         if len(history_images) >= history_steps:
                             history_images = history_images[-history_steps + 1 :]
+                            history_depths = history_depths[-history_steps + 1 :]
+                            history_depth_valid = history_depth_valid[-history_steps + 1 :]
                             history_proprio = history_proprio[-history_steps + 1 :]
                             history_actions = history_actions[-history_steps + 1 :]
                         history_images.append(observation["images"])
+                        if "depths" in observation:
+                            history_depths.append(observation["depths"])
+                            history_depth_valid.append(observation["depth_valid"])
                         history_proprio.append(observation["proprio"])
                         history_actions.append(action.astype(np.float32))
                     observation, _, terminated, truncated, privileged_state = env.step(action)

code/reveal_vla_bimanual/eval/run_rlbench_rollout_eval.py

@@ -52,6 +52,19 @@ def _episode_language_goal(descriptions: Sequence[str]) -> str:
     return str(descriptions[0]) if descriptions else ""
 
 
+def _reset_task_with_retries(task: Any, max_attempts: int) -> tuple[Sequence[str], Any, int]:
+    last_error: Exception | None = None
+    for attempt in range(max_attempts):
+        try:
+            descriptions, obs = task.reset()
+            return descriptions, obs, attempt
+        except Exception as exc:  # pragma: no cover - live RLBench failure path
+            last_error = exc
+    if last_error is not None:
+        raise last_error
+    raise RuntimeError("Task reset failed without raising a concrete exception.")
+
+
 def _noop_bimanual_action(obs: Any) -> np.ndarray:
     right_obs = getattr(obs, "right", None)
     left_obs = getattr(obs, "left", None)
@@ -113,6 +126,7 @@ def main() -> None:
     parser.add_argument("--disable-support-mode-conditioning", action="store_true")
     parser.add_argument("--headless", action="store_true", default=True)
     parser.add_argument("--chunk-commit-steps", type=int, default=0)
+    parser.add_argument("--reset-retries", type=int, default=20)
     args = parser.parse_args()
 
     checkpoint = torch.load(Path(args.checkpoint), map_location="cpu", weights_only=False)
@@ -155,6 +169,7 @@ def main() -> None:
         "episodes_per_task": args.episodes_per_task,
         "episode_length": args.episode_length,
         "resolution": args.resolution,
+        "reset_retries": args.reset_retries,
         "cameras": list(camera_spec.cameras),
         "tasks": {},
     }
@@ -180,8 +195,10 @@ def main() -> None:
             )
             env.launch()
             task = env.get_task(task_class)
+            task_reset_retries: list[int] = []
             for _ in range(args.episodes_per_task):
-                descriptions, obs = task.reset()
+                descriptions, obs, reset_retries = _reset_task_with_retries(task, max_attempts=max(1, args.reset_retries))
+                task_reset_retries.append(int(reset_retries))
                 language_goal = _episode_language_goal(descriptions)
                 total_reward = 0.0
                 success = 0.0
@@ -291,6 +308,7 @@ def main() -> None:
                 "returns": task_returns,
                 "path_recoveries": episode_recoveries if args.episodes_per_task == 1 else None,
                 "noop_fallbacks": episode_noop_fallbacks if args.episodes_per_task == 1 else None,
+                "reset_retries": task_reset_retries,
                 "mean_success": float(np.mean(task_successes)) if task_successes else 0.0,
                 "mean_return": float(np.mean(task_returns)) if task_returns else 0.0,
             }

code/reveal_vla_bimanual/models/action_decoder.py

@@ -19,6 +19,8 @@ class ChunkDecoderConfig:
     num_candidates: int = 8
     num_phases: int = 5
     num_arm_roles: int = 4
+    num_proposal_modes: int = 6
+    planner_top_k: int = 4
 
 
 class ACTBimanualChunkDecoder(nn.Module):
@@ -381,3 +383,305 @@ class InteractionChunkDecoder(nn.Module):
         candidates = action_mean.unsqueeze(1) + noise * action_log_std.exp().unsqueeze(1)
         candidates[:, 0] = action_mean
         return candidates
+
+
+DEFAULT_PROPOSAL_MODES = (
+    "widen_opening",
+    "maintain_opening",
+    "slide_occluder",
+    "lift_support_layer",
+    "stabilize_support",
+    "retrieve",
+)
+
+
+def swap_arm_action_order(action_chunk: Tensor) -> Tensor:
+    midpoint = action_chunk.shape[-1] // 2
+    return torch.cat([action_chunk[..., midpoint:], action_chunk[..., :midpoint]], dim=-1)
+
+
+class SymmetricCoordinatedChunkDecoder(nn.Module):
+    def __init__(self, config: ChunkDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        proposal_context_dim = config.action_dim + (config.hidden_dim * 2)
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=config.hidden_dim,
+            nhead=config.num_heads,
+            dim_feedforward=config.ff_dim,
+            dropout=config.dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.arm_decoder = nn.TransformerDecoder(decoder_layer, num_layers=config.num_layers)
+        self.query_embed = nn.Embedding(config.chunk_size, config.hidden_dim)
+        self.arm_identity = nn.Embedding(2, config.hidden_dim)
+        self.phase_adapter = nn.Linear(config.num_phases, config.hidden_dim)
+        self.role_adapter = nn.Linear(config.num_arm_roles, config.hidden_dim)
+        self.context_proj = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.coordination = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim * 3),
+            nn.Linear(config.hidden_dim * 3, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+        )
+        self.arm_head = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.arm_mean = nn.Linear(config.hidden_dim, config.arm_action_dim)
+        self.arm_log_std = nn.Linear(config.hidden_dim, config.arm_action_dim)
+        self.proposal_mode_head = nn.Sequential(
+            nn.LayerNorm(proposal_context_dim),
+            nn.Linear(proposal_context_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.num_proposal_modes),
+        )
+        self.proposal_mode_embeddings = nn.Embedding(config.num_proposal_modes, config.hidden_dim)
+        self.proposal_slot_embeddings = nn.Embedding(config.num_candidates, config.hidden_dim)
+        self.mode_residual_heads = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.LayerNorm(proposal_context_dim),
+                    nn.Linear(proposal_context_dim, config.hidden_dim),
+                    nn.GELU(),
+                    nn.Linear(config.hidden_dim, config.chunk_size * config.action_dim),
+                )
+                for _ in range(config.num_proposal_modes)
+            ]
+        )
+        self.slot_delta = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.chunk_size * config.action_dim),
+        )
+        self.proposal_score = nn.Sequential(
+            nn.LayerNorm(proposal_context_dim + config.hidden_dim),
+            nn.Linear(proposal_context_dim + config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, 1),
+        )
+
+    def _conditioning(
+        self,
+        interaction_state: dict[str, Tensor] | None,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        swap_roles: bool = False,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        if interaction_state is None:
+            zero_phase = torch.zeros(batch_size, self.config.hidden_dim, device=device, dtype=dtype)
+            zero_roles = torch.zeros(batch_size, 2, self.config.hidden_dim, device=device, dtype=dtype)
+            zero_context = torch.zeros(batch_size, self.config.hidden_dim, device=device, dtype=dtype)
+            return zero_phase, zero_roles, zero_context
+        phase_probs = interaction_state["phase_logits"].softmax(dim=-1).to(dtype=dtype)
+        arm_role_probs = interaction_state["arm_role_logits"].softmax(dim=-1).to(dtype=dtype)
+        if swap_roles:
+            arm_role_probs = arm_role_probs.flip(1)
+        phase_context = self.phase_adapter(phase_probs)
+        role_context = self.role_adapter(arm_role_probs)
+        if interaction_state.get("interaction_tokens") is not None:
+            interaction_context = interaction_state["interaction_tokens"].mean(dim=1)
+        else:
+            interaction_context = interaction_state["field_tokens"].mean(dim=1)
+        return phase_context, role_context, self.context_proj(interaction_context)
+
+    def _decode_arm_tokens(
+        self,
+        queries: Tensor,
+        decoder_memory: Tensor,
+        phase_context: Tensor,
+        role_context: Tensor,
+        interaction_context: Tensor,
+        swap_roles: bool = False,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        batch_size, chunk_size, _ = queries.shape
+        identity_order = torch.tensor([1, 0], device=queries.device) if swap_roles else torch.tensor([0, 1], device=queries.device)
+        arm_queries = queries.unsqueeze(1).expand(-1, 2, -1, -1)
+        arm_queries = arm_queries + phase_context.unsqueeze(1).unsqueeze(2)
+        arm_queries = arm_queries + role_context.unsqueeze(2)
+        arm_queries = arm_queries + self.arm_identity(identity_order).view(1, 2, 1, -1).to(dtype=queries.dtype)
+        flat_queries = arm_queries.reshape(batch_size * 2, chunk_size, self.config.hidden_dim)
+        flat_memory = decoder_memory.unsqueeze(1).expand(-1, 2, -1, -1).reshape(
+            batch_size * 2,
+            decoder_memory.shape[1],
+            decoder_memory.shape[2],
+        )
+        decoded = self.arm_decoder(flat_queries, flat_memory).reshape(batch_size, 2, chunk_size, self.config.hidden_dim)
+        coordination_input = torch.cat(
+            [
+                decoded[:, 0],
+                decoded[:, 1],
+                interaction_context.unsqueeze(1).expand(-1, chunk_size, -1),
+            ],
+            dim=-1,
+        )
+        coordination = torch.tanh(self.coordination(coordination_input))
+        decoded[:, 0] = decoded[:, 0] + coordination
+        decoded[:, 1] = decoded[:, 1] + coordination
+        decoded = self.arm_head(decoded)
+        arm_mean = self.arm_mean(decoded)
+        arm_log_std = self.arm_log_std(decoded).clamp(min=-5.0, max=2.0)
+        return arm_mean, arm_log_std, coordination
+
+    def _proposal_outputs(
+        self,
+        base_action: Tensor,
+        pooled_context: Tensor,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        batch_size = pooled_context.shape[0]
+        mode_logits = self.proposal_mode_head(pooled_context)
+        mode_residuals = []
+        for head in self.mode_residual_heads:
+            residual = head(pooled_context).view(batch_size, self.config.chunk_size, self.config.action_dim)
+            mode_residuals.append(residual)
+        mode_residuals = torch.stack(mode_residuals, dim=1)
+
+        mode_assignments = torch.arange(self.config.num_candidates, device=pooled_context.device) % self.config.num_proposal_modes
+        slot_embeddings = self.proposal_slot_embeddings.weight
+        slot_deltas = self.slot_delta(slot_embeddings).view(
+            self.config.num_candidates,
+            self.config.chunk_size,
+            self.config.action_dim,
+        )
+        proposal_candidates = []
+        proposal_logits = []
+        for slot_idx in range(self.config.num_candidates):
+            mode_idx = int(mode_assignments[slot_idx])
+            candidate = base_action + 0.35 * torch.tanh(mode_residuals[:, mode_idx]) + 0.05 * torch.tanh(slot_deltas[slot_idx]).unsqueeze(0)
+            proposal_candidates.append(candidate)
+            score_features = torch.cat(
+                [
+                    pooled_context,
+                    self.proposal_mode_embeddings.weight[mode_idx].unsqueeze(0).expand(batch_size, -1)
+                    + slot_embeddings[slot_idx].unsqueeze(0).expand(batch_size, -1),
+                ],
+                dim=-1,
+            )
+            proposal_logits.append(
+                self.proposal_score(score_features).squeeze(-1) + mode_logits[:, mode_idx]
+            )
+        stacked_candidates = torch.stack(proposal_candidates, dim=1)
+        stacked_logits = torch.stack(proposal_logits, dim=1)
+        stacked_candidates[:, 0] = base_action
+        return stacked_candidates, stacked_logits, mode_logits
+
+    def forward(
+        self,
+        scene_tokens: Tensor,
+        interaction_state: dict[str, Tensor] | None = None,
+        memory_tokens: Tensor | None = None,
+        reveal_tokens: Tensor | None = None,
+        memory_token: Tensor | None = None,
+        compute_equivariance_probe: bool = False,
+    ) -> dict[str, Tensor]:
+        if memory_tokens is None:
+            memory_tokens = memory_token
+        batch_size = scene_tokens.shape[0]
+        dtype = scene_tokens.dtype
+        phase_context, role_context, interaction_context = self._conditioning(
+            interaction_state=interaction_state,
+            batch_size=batch_size,
+            device=scene_tokens.device,
+            dtype=dtype,
+        )
+
+        decoder_memory = scene_tokens
+        interaction_tokens = interaction_state.get("interaction_tokens") if interaction_state is not None else None
+        if interaction_tokens is not None:
+            decoder_memory = torch.cat([decoder_memory, interaction_tokens], dim=1)
+        elif reveal_tokens is not None:
+            decoder_memory = torch.cat([decoder_memory, reveal_tokens], dim=1)
+        if memory_tokens is not None:
+            decoder_memory = torch.cat([decoder_memory, memory_tokens], dim=1)
+
+        base_queries = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        arm_mean, arm_log_std, coordination = self._decode_arm_tokens(
+            queries=base_queries,
+            decoder_memory=decoder_memory,
+            phase_context=phase_context,
+            role_context=role_context,
+            interaction_context=interaction_context,
+        )
+        action_mean = torch.cat([arm_mean[:, 0], arm_mean[:, 1]], dim=-1)
+        action_log_std = torch.cat([arm_log_std[:, 0], arm_log_std[:, 1]], dim=-1)
+        pooled_context = torch.cat(
+            [
+                arm_mean[:, 0].mean(dim=1),
+                arm_mean[:, 1].mean(dim=1),
+                coordination.mean(dim=1),
+                interaction_context,
+            ],
+            dim=-1,
+        )
+        proposal_candidates, proposal_logits, proposal_mode_logits = self._proposal_outputs(action_mean, pooled_context)
+
+        outputs = {
+            "decoded_tokens": torch.cat([arm_mean[:, 0], arm_mean[:, 1]], dim=-1),
+            "right_tokens": arm_mean[:, 0],
+            "left_tokens": arm_mean[:, 1],
+            "coordination_tokens": coordination,
+            "action_mean": action_mean,
+            "action_log_std": action_log_std,
+            "proposal_candidates": proposal_candidates,
+            "proposal_logits": proposal_logits,
+            "proposal_mode_logits": proposal_mode_logits,
+            "proposal_mode_assignments": torch.arange(
+                self.config.num_candidates,
+                device=scene_tokens.device,
+            ) % self.config.num_proposal_modes,
+            "proposal_mode_names": list(DEFAULT_PROPOSAL_MODES[: self.config.num_proposal_modes]),
+        }
+        if compute_equivariance_probe:
+            swapped_phase, swapped_roles, swapped_context = self._conditioning(
+                interaction_state=interaction_state,
+                batch_size=batch_size,
+                device=scene_tokens.device,
+                dtype=dtype,
+                swap_roles=True,
+            )
+            swapped_arm_mean, _, _ = self._decode_arm_tokens(
+                queries=base_queries,
+                decoder_memory=decoder_memory,
+                phase_context=swapped_phase,
+                role_context=swapped_roles,
+                interaction_context=swapped_context,
+                swap_roles=True,
+            )
+            outputs["equivariance_probe_action_mean"] = torch.cat(
+                [swapped_arm_mean[:, 0], swapped_arm_mean[:, 1]],
+                dim=-1,
+            )
+            outputs["equivariance_target_action_mean"] = swap_arm_action_order(action_mean)
+        return outputs
+
+    def sample_candidates(
+        self,
+        action_mean: Tensor,
+        action_log_std: Tensor,
+        num_candidates: int | None = None,
+        proposal_candidates: Tensor | None = None,
+    ) -> Tensor:
+        if proposal_candidates is not None:
+            return proposal_candidates
+        num_candidates = num_candidates or self.config.num_candidates
+        if num_candidates <= 1:
+            return action_mean.unsqueeze(1)
+        noise = torch.randn(
+            action_mean.size(0),
+            num_candidates,
+            action_mean.size(1),
+            action_mean.size(2),
+            device=action_mean.device,
+            dtype=action_mean.dtype,
+        )
+        candidates = action_mean.unsqueeze(1) + noise * action_log_std.exp().unsqueeze(1)
+        candidates[:, 0] = action_mean
+        return candidates

code/reveal_vla_bimanual/models/backbones.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 import math
+import os
 from pathlib import Path
 from typing import Sequence
 
@@ -18,6 +19,157 @@ class FrozenVLBackboneConfig:
     freeze_backbone: bool = True
     gradient_checkpointing: bool = True
     use_dummy_backbone: bool = False
+    depth_patch_size: int = 16
+    geometry_feature_dim: int = 8
+    use_camera_geometry: bool = True
+
+
+class DepthPatchAdapter(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        patch_size: int = 16,
+        geometry_feature_dim: int = 8,
+    ) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.patch_size = patch_size
+        self.geometry_feature_dim = geometry_feature_dim
+        self.depth_proj = nn.Sequential(
+            nn.LayerNorm(2 + geometry_feature_dim),
+            nn.Linear(2 + geometry_feature_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+        self.geometry_proj = nn.Sequential(
+            nn.LayerNorm(geometry_feature_dim),
+            nn.Linear(geometry_feature_dim, hidden_dim),
+            nn.GELU(),
+        )
+        self.camera_proj = nn.Sequential(
+            nn.LayerNorm(7),
+            nn.Linear(7, hidden_dim),
+            nn.GELU(),
+        )
+
+    def _patchify(self, tensor: Tensor) -> Tensor:
+        pooled = F.avg_pool2d(tensor, kernel_size=self.patch_size, stride=self.patch_size)
+        return pooled.flatten(2).transpose(1, 2)
+
+    def _geometry_features(
+        self,
+        depths: Tensor,
+        camera_intrinsics: Tensor | None = None,
+        camera_extrinsics: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        batch_views, _, height, width = depths.shape
+        grid_h = max(1, height // self.patch_size)
+        grid_w = max(1, width // self.patch_size)
+        y_coords = torch.linspace(-1.0, 1.0, steps=grid_h, device=depths.device, dtype=depths.dtype)
+        x_coords = torch.linspace(-1.0, 1.0, steps=grid_w, device=depths.device, dtype=depths.dtype)
+        grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
+        coords = torch.stack([grid_x, grid_y], dim=-1).reshape(1, grid_h * grid_w, 2).expand(batch_views, -1, -1)
+
+        geometry_terms: list[Tensor] = [coords]
+        if camera_intrinsics is not None:
+            fx = camera_intrinsics[:, 0, 0].unsqueeze(-1).unsqueeze(-1)
+            fy = camera_intrinsics[:, 1, 1].unsqueeze(-1).unsqueeze(-1)
+            cx = camera_intrinsics[:, 0, 2].unsqueeze(-1).unsqueeze(-1)
+            cy = camera_intrinsics[:, 1, 2].unsqueeze(-1).unsqueeze(-1)
+            intrinsic_features = torch.cat(
+                [
+                    fx.expand(-1, grid_h * grid_w, -1),
+                    fy.expand(-1, grid_h * grid_w, -1),
+                    cx.expand(-1, grid_h * grid_w, -1),
+                    cy.expand(-1, grid_h * grid_w, -1),
+                ],
+                dim=-1,
+            )
+            geometry_terms.append(intrinsic_features)
+        else:
+            geometry_terms.append(torch.zeros(batch_views, grid_h * grid_w, 4, device=depths.device, dtype=depths.dtype))
+
+        if camera_extrinsics is not None:
+            translation = camera_extrinsics[:, :3, 3]
+            translation = translation.unsqueeze(1).expand(-1, grid_h * grid_w, -1)
+            geometry_terms.append(translation)
+        else:
+            geometry_terms.append(torch.zeros(batch_views, grid_h * grid_w, 3, device=depths.device, dtype=depths.dtype))
+
+        geometry = torch.cat(geometry_terms, dim=-1)
+        if geometry.shape[-1] < self.geometry_feature_dim:
+            pad = self.geometry_feature_dim - geometry.shape[-1]
+            geometry = F.pad(geometry, (0, pad))
+        elif geometry.shape[-1] > self.geometry_feature_dim:
+            geometry = geometry[..., : self.geometry_feature_dim]
+
+        if camera_intrinsics is not None:
+            camera_summary = torch.cat(
+                [
+                    camera_intrinsics[:, 0, 0:1],
+                    camera_intrinsics[:, 1, 1:2],
+                    camera_intrinsics[:, 0, 2:3],
+                    camera_intrinsics[:, 1, 2:3],
+                ],
+                dim=-1,
+            )
+        else:
+            camera_summary = torch.zeros(batch_views, 4, device=depths.device, dtype=depths.dtype)
+        if camera_extrinsics is not None:
+            camera_summary = torch.cat([camera_summary, camera_extrinsics[:, :3, 3]], dim=-1)
+        else:
+            camera_summary = torch.cat(
+                [camera_summary, torch.zeros(batch_views, 3, device=depths.device, dtype=depths.dtype)],
+                dim=-1,
+            )
+        return geometry, camera_summary
+
+    def forward(
+        self,
+        depths: Tensor,
+        depth_valid: Tensor | None = None,
+        camera_intrinsics: Tensor | None = None,
+        camera_extrinsics: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        if depths.ndim == 4:
+            depths = depths.unsqueeze(2)
+        if depth_valid is None:
+            depth_valid = torch.ones_like(depths)
+        if depth_valid.ndim == 4:
+            depth_valid = depth_valid.unsqueeze(2)
+        if depths.ndim != 5:
+            raise ValueError(f"Expected depths to have shape [B, V, H, W] or [B, V, 1, H, W], got {tuple(depths.shape)}")
+        if depths.shape[2] != 1:
+            depths = depths.mean(dim=2, keepdim=True)
+        if depth_valid.shape[2] != 1:
+            depth_valid = depth_valid.mean(dim=2, keepdim=True)
+
+        batch_size, num_views = depths.shape[:2]
+        flat_depths = depths.reshape(batch_size * num_views, 1, depths.shape[-2], depths.shape[-1]).float()
+        flat_valid = depth_valid.reshape(batch_size * num_views, 1, depth_valid.shape[-2], depth_valid.shape[-1]).float()
+        flat_intrinsics = None
+        flat_extrinsics = None
+        if camera_intrinsics is not None:
+            flat_intrinsics = camera_intrinsics.reshape(batch_size * num_views, *camera_intrinsics.shape[-2:]).float()
+        if camera_extrinsics is not None:
+            flat_extrinsics = camera_extrinsics.reshape(batch_size * num_views, *camera_extrinsics.shape[-2:]).float()
+
+        depth_patch = self._patchify(flat_depths)
+        valid_patch = self._patchify(flat_valid)
+        geometry_features, camera_summary = self._geometry_features(
+            flat_depths,
+            camera_intrinsics=flat_intrinsics,
+            camera_extrinsics=flat_extrinsics,
+        )
+        token_inputs = torch.cat([depth_patch, valid_patch, geometry_features], dim=-1)
+        depth_tokens = self.depth_proj(token_inputs)
+        geometry_tokens = self.geometry_proj(geometry_features)
+        camera_tokens = self.camera_proj(camera_summary).unsqueeze(1)
+        return {
+            "depth_tokens": depth_tokens.view(batch_size, num_views, depth_tokens.shape[1], depth_tokens.shape[2]),
+            "geometry_tokens": geometry_tokens.view(batch_size, num_views, geometry_tokens.shape[1], geometry_tokens.shape[2]),
+            "camera_tokens": camera_tokens.view(batch_size, num_views, 1, camera_tokens.shape[-1]),
+        }
 
 
 class _DummyTextTokenizer:
@@ -42,6 +194,11 @@ class FrozenVLBackbone(nn.Module):
         self.config = config
         self.hidden_dim = config.hidden_dim
         self.use_dummy_backbone = config.use_dummy_backbone
+        self.depth_adapter = DepthPatchAdapter(
+            hidden_dim=config.hidden_dim,
+            patch_size=config.depth_patch_size,
+            geometry_feature_dim=config.geometry_feature_dim,
+        )
 
         if config.use_dummy_backbone:
             self.image_patch_size = 16
@@ -51,36 +208,62 @@ class FrozenVLBackbone(nn.Module):
 
             local_model_source: str | None = None
             if config.model_name == "openai/clip-vit-base-patch32":
-                cache_root = Path("/workspace/.cache/huggingface/hub/models--openai--clip-vit-base-patch32")
-                ref_path = cache_root / "refs" / "main"
-                if ref_path.exists():
-                    snapshot_id = ref_path.read_text(encoding="utf-8").strip()
-                    snapshot_dir = cache_root / "snapshots" / snapshot_id
-                    if (snapshot_dir / "config.json").exists():
-                        local_model_source = str(snapshot_dir)
+                explicit_local_dir = Path("/workspace/models/openai_clip_vit_base_patch32")
+                if (explicit_local_dir / "config.json").exists():
+                    local_model_source = str(explicit_local_dir)
+                cache_home = Path(os.environ.get("HF_HOME", "/workspace/.cache/huggingface"))
+                cache_root = cache_home / "hub" / "models--openai--clip-vit-base-patch32"
+                if local_model_source is None:
+                    ref_path = cache_root / "refs" / "main"
+                    if ref_path.exists():
+                        snapshot_id = ref_path.read_text(encoding="utf-8").strip()
+                        snapshot_dir = cache_root / "snapshots" / snapshot_id
+                        if (snapshot_dir / "config.json").exists():
+                            local_model_source = str(snapshot_dir)
+                if local_model_source is None:
+                    snapshot_root = cache_root / "snapshots"
+                    if snapshot_root.exists():
+                        for snapshot_dir in sorted(snapshot_root.iterdir(), reverse=True):
+                            if (snapshot_dir / "config.json").exists():
+                                local_model_source = str(snapshot_dir)
+                                break
             clip_model = None
+            last_clip_error: Exception | None = None
+            model_sources: list[tuple[str, dict[str, object]]] = []
             if local_model_source is not None:
+                model_sources.append((local_model_source, {"use_safetensors": True, "local_files_only": True}))
+                model_sources.append((local_model_source, {"local_files_only": True}))
+            model_sources.append((config.model_name, {"use_safetensors": True}))
+            model_sources.append((config.model_name, {}))
+            for source, kwargs in model_sources:
                 try:
-                    clip_model = CLIPModel.from_pretrained(
-                        local_model_source,
-                        use_safetensors=True,
-                        local_files_only=True,
-                    )
-                except OSError:
-                    clip_model = None
+                    clip_model = CLIPModel.from_pretrained(source, **kwargs)
+                    break
+                except Exception as exc:
+                    last_clip_error = exc
             if clip_model is None:
-                clip_model = CLIPModel.from_pretrained(config.model_name, use_safetensors=True)
+                assert last_clip_error is not None
+                raise last_clip_error
             self.vision_model = clip_model.vision_model
             self.text_model = clip_model.text_model
             self.visual_projection = clip_model.visual_projection
             self.text_projection = clip_model.text_projection
+            tokenizer = None
+            last_tokenizer_error: Exception | None = None
+            tokenizer_sources: list[tuple[str, dict[str, object]]] = []
             if local_model_source is not None:
+                tokenizer_sources.append((local_model_source, {"local_files_only": True}))
+            tokenizer_sources.append((config.model_name, {}))
+            for source, kwargs in tokenizer_sources:
                 try:
-                    self.tokenizer = AutoTokenizer.from_pretrained(local_model_source, local_files_only=True)
-                except OSError:
-                    self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
-            else:
-                self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+                    tokenizer = AutoTokenizer.from_pretrained(source, **kwargs)
+                    break
+                except Exception as exc:
+                    last_tokenizer_error = exc
+            if tokenizer is None:
+                assert last_tokenizer_error is not None
+                raise last_tokenizer_error
+            self.tokenizer = tokenizer
             self.hidden_dim = clip_model.config.projection_dim
             if config.gradient_checkpointing:
                 if hasattr(self.vision_model, "gradient_checkpointing_enable"):
@@ -88,9 +271,17 @@ class FrozenVLBackbone(nn.Module):
                 if hasattr(self.text_model, "gradient_checkpointing_enable"):
                     self.text_model.gradient_checkpointing_enable()
 
-        if config.freeze_backbone:
-            for parameter in self.parameters():
-                parameter.requires_grad = False
+        if config.freeze_backbone and not config.use_dummy_backbone:
+            for module in (
+                getattr(self, "vision_model", None),
+                getattr(self, "text_model", None),
+                getattr(self, "visual_projection", None),
+                getattr(self, "text_projection", None),
+            ):
+                if module is None:
+                    continue
+                for parameter in module.parameters():
+                    parameter.requires_grad = False
 
     def tokenize_text(self, texts: Sequence[str], device: torch.device) -> dict[str, Tensor]:
         if self.use_dummy_backbone:
@@ -103,7 +294,7 @@ class FrozenVLBackbone(nn.Module):
             return_tensors="pt",
         ).to(device)
 
-    def encode_images(self, images: Tensor) -> Tensor:
+    def _encode_rgb_tokens(self, images: Tensor) -> Tensor:
         batch_size, num_views, channels, height, width = images.shape
         flat_images = images.reshape(batch_size * num_views, channels, height, width)
         if self.use_dummy_backbone:
@@ -125,6 +316,40 @@ class FrozenVLBackbone(nn.Module):
         num_tokens = tokens.shape[1]
         return tokens.reshape(batch_size, num_views, num_tokens, -1)
 
+    def encode_images(
+        self,
+        images: Tensor,
+        depths: Tensor | None = None,
+        depth_valid: Tensor | None = None,
+        camera_intrinsics: Tensor | None = None,
+        camera_extrinsics: Tensor | None = None,
+        return_aux: bool = False,
+    ) -> Tensor | dict[str, Tensor | None]:
+        rgb_tokens = self._encode_rgb_tokens(images)
+        wants_aux = return_aux or depths is not None or depth_valid is not None or camera_intrinsics is not None or camera_extrinsics is not None
+        if not wants_aux:
+            return rgb_tokens
+
+        depth_outputs: dict[str, Tensor | None] = {
+            "depth_tokens": None,
+            "geometry_tokens": None,
+            "camera_tokens": None,
+        }
+        if depths is not None:
+            depth_outputs = self.depth_adapter(
+                depths=depths,
+                depth_valid=depth_valid,
+                camera_intrinsics=camera_intrinsics,
+                camera_extrinsics=camera_extrinsics,
+            )
+
+        return {
+            "rgb_tokens": rgb_tokens,
+            "depth_tokens": depth_outputs["depth_tokens"],
+            "geometry_tokens": depth_outputs["geometry_tokens"],
+            "camera_tokens": depth_outputs["camera_tokens"],
+        }
+
     def encode_text(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
         if self.use_dummy_backbone:
             vocab_scale = float(self.tokenizer.vocab_size - 1)

code/reveal_vla_bimanual/models/multiview_fusion.py

@@ -16,6 +16,37 @@ class MultiViewFusionConfig:
     dropout: float = 0.1
     proprio_dim: int = 32
     proprio_tokens: int = 1
+    geometry_num_heads: int = 4
+
+
+class GatedCrossAttentionBlock(nn.Module):
+    def __init__(self, hidden_dim: int, num_heads: int, dropout: float) -> None:
+        super().__init__()
+        self.attn = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.gate = nn.Sequential(
+            nn.LayerNorm(hidden_dim * 2),
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+        self.out = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+        )
+
+    def forward(self, rgb_tokens: Tensor, geometry_tokens: Tensor) -> tuple[Tensor, Tensor]:
+        attended, _ = self.attn(rgb_tokens, geometry_tokens, geometry_tokens)
+        rgb_summary = rgb_tokens.mean(dim=1)
+        geometry_summary = geometry_tokens.mean(dim=1)
+        gate = torch.sigmoid(self.gate(torch.cat([rgb_summary, geometry_summary], dim=-1))).unsqueeze(1)
+        fused = rgb_tokens + gate * attended
+        return self.out(fused), geometry_summary
 
 
 class MultiViewFusion(nn.Module):
@@ -35,13 +66,26 @@ class MultiViewFusion(nn.Module):
             encoder_layer,
             num_layers=config.num_layers,
         )
+        self.geometry_fusion = GatedCrossAttentionBlock(
+            hidden_dim=config.hidden_dim,
+            num_heads=max(1, min(config.num_heads, config.geometry_num_heads)),
+            dropout=config.dropout,
+        )
         self.proprio_adapter = nn.Sequential(
             nn.LayerNorm(config.proprio_dim),
             nn.Linear(config.proprio_dim, config.hidden_dim * config.proprio_tokens),
             nn.GELU(),
         )
 
-    def forward(self, image_tokens: Tensor, proprio: Tensor, language_tokens: Tensor) -> Tensor:
+    def forward(
+        self,
+        image_tokens: Tensor,
+        proprio: Tensor,
+        language_tokens: Tensor,
+        depth_tokens: Tensor | None = None,
+        camera_tokens: Tensor | None = None,
+        return_aux: bool = False,
+    ) -> Tensor | dict[str, Tensor]:
         batch_size, num_views, num_tokens, hidden_dim = image_tokens.shape
         if num_views != self.config.num_cameras:
             raise ValueError(f"Expected {self.config.num_cameras} views, received {num_views}")
@@ -49,9 +93,36 @@ class MultiViewFusion(nn.Module):
         camera_ids = torch.arange(num_views, device=image_tokens.device)
         camera_embed = self.camera_embedding(camera_ids).view(1, num_views, 1, hidden_dim)
         image_tokens = image_tokens + camera_embed
-        fused = self.cross_view_transformer(image_tokens.reshape(batch_size, num_views * num_tokens, hidden_dim))
+
+        per_view_tokens = []
+        view_summaries = []
+        geometry_summaries = []
+        for view_idx in range(num_views):
+            rgb_tokens = image_tokens[:, view_idx]
+            geometry_sources = []
+            if depth_tokens is not None:
+                geometry_sources.append(depth_tokens[:, view_idx])
+            if camera_tokens is not None:
+                geometry_sources.append(camera_tokens[:, view_idx])
+            if geometry_sources:
+                geometry = torch.cat(geometry_sources, dim=1)
+                rgb_tokens, geometry_summary = self.geometry_fusion(rgb_tokens, geometry)
+            else:
+                geometry_summary = torch.zeros(batch_size, hidden_dim, device=image_tokens.device, dtype=image_tokens.dtype)
+            per_view_tokens.append(rgb_tokens)
+            view_summaries.append(rgb_tokens.mean(dim=1))
+            geometry_summaries.append(geometry_summary)
+
+        fused = self.cross_view_transformer(torch.cat(per_view_tokens, dim=1))
 
         proprio_tokens = self.proprio_adapter(proprio).view(
             batch_size, self.config.proprio_tokens, hidden_dim
         )
-        return torch.cat([fused, proprio_tokens, language_tokens], dim=1)
+        scene_tokens = torch.cat([fused, proprio_tokens, language_tokens], dim=1)
+        if not (return_aux or depth_tokens is not None or camera_tokens is not None):
+            return scene_tokens
+        return {
+            "scene_tokens": scene_tokens,
+            "view_summaries": torch.stack(view_summaries, dim=1),
+            "geometry_summaries": torch.stack(geometry_summaries, dim=1),
+        }

code/reveal_vla_bimanual/models/observation_memory.py

@@ -16,6 +16,12 @@ class ObservationMemoryConfig:
     memory_bank_size: int = 4
     num_heads: int = 4
     max_history_steps: int = 8
+    scene_bank_size: int = 2
+    belief_bank_size: int = 2
+    scene_history_steps: int = 3
+    belief_history_steps: int = 8
+    memory_write_threshold: float = 0.45
+    memory_suppression_margin: float = 0.05
 
 
 class ObservationMemory(nn.Module):
@@ -173,3 +179,189 @@ class InteractionObservationMemory(nn.Module):
             "memory_tokens": projected_bank,
             "memory_uncertainty": torch.nn.functional.softplus(self.uncertainty_head(pooled_bank)).squeeze(-1),
         }
+
+
+class _SelectiveMemoryBank(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        action_dim: int,
+        num_heads: int,
+        dropout: float,
+        bank_size: int,
+        history_steps: int,
+        max_history_steps: int,
+        write_threshold: float,
+        suppression_margin: float,
+    ) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.history_steps = history_steps
+        self.max_history_steps = max_history_steps
+        self.write_threshold = write_threshold
+        self.suppression_margin = suppression_margin
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_dim * 4,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.sequence_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
+        self.position_embedding = nn.Parameter(torch.randn(1, max_history_steps + 1, hidden_dim) * 0.02)
+        self.bank_queries = nn.Parameter(torch.randn(bank_size, hidden_dim) * 0.02)
+        self.bank_attention = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.action_proj = nn.Sequential(
+            nn.LayerNorm(action_dim),
+            nn.Linear(action_dim, hidden_dim),
+            nn.GELU(),
+        )
+        self.write_gate = nn.Sequential(
+            nn.LayerNorm(hidden_dim * 3),
+            nn.Linear(hidden_dim * 3, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, 1),
+        )
+        self.token_proj = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+        )
+
+    def _truncate(self, history: Tensor | None) -> Tensor | None:
+        if history is None or history.numel() == 0:
+            return history
+        if history.shape[1] <= self.history_steps:
+            return history
+        return history[:, -self.history_steps :]
+
+    def forward(
+        self,
+        pooled_current: Tensor,
+        history_scene_tokens: Tensor | None = None,
+        history_actions: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        history_scene_tokens = self._truncate(history_scene_tokens)
+        pooled_current = pooled_current.unsqueeze(1)
+        if history_scene_tokens is not None and history_scene_tokens.numel() > 0:
+            history_pooled = history_scene_tokens.mean(dim=2)
+            if history_actions is not None and history_actions.numel() > 0:
+                history_actions = history_actions[:, -history_pooled.shape[1] :]
+                history_pooled = history_pooled + self.action_proj(history_actions)
+            sequence = torch.cat([history_pooled, pooled_current], dim=1)
+        else:
+            history_pooled = pooled_current[:, :0]
+            sequence = pooled_current
+        if sequence.shape[1] > self.position_embedding.shape[1]:
+            raise ValueError(
+                f"Sequence length {sequence.shape[1]} exceeds configured maximum {self.position_embedding.shape[1]}"
+            )
+        encoded = self.sequence_encoder(sequence + self.position_embedding[:, : sequence.shape[1]])
+        current_token = encoded[:, -1]
+        prior_token = encoded[:, :-1].mean(dim=1) if encoded.shape[1] > 1 else torch.zeros_like(current_token)
+        novelty = torch.abs(current_token - prior_token)
+        informative = novelty.mean(dim=-1, keepdim=True)
+        gate_logit = self.write_gate(torch.cat([current_token, prior_token, novelty], dim=-1))
+        gate = torch.sigmoid(gate_logit)
+        gate = gate * (informative > (self.write_threshold - self.suppression_margin)).to(gate.dtype)
+        recent_summary = encoded[:, -min(max(1, self.bank_queries.shape[0]), encoded.shape[1]) :].mean(dim=1, keepdim=True)
+        queries = self.bank_queries.unsqueeze(0).expand(encoded.shape[0], -1, -1) + recent_summary
+        bank_tokens, _ = self.bank_attention(queries, encoded, encoded)
+        bank_tokens = bank_tokens + recent_summary
+        bank_tokens = prior_token.unsqueeze(1) * (1.0 - gate.unsqueeze(1)) + bank_tokens * gate.unsqueeze(1)
+        bank_tokens = self.token_proj(bank_tokens)
+        return {
+            "memory_tokens": bank_tokens,
+            "memory_token": bank_tokens.mean(dim=1, keepdim=True),
+            "memory_sequence": encoded,
+            "memory_state": current_token,
+            "write_gate": gate.squeeze(-1),
+            "saturation": bank_tokens.abs().mean(dim=(1, 2)),
+        }
+
+
+class SceneMemory(_SelectiveMemoryBank):
+    def __init__(self, config: ObservationMemoryConfig) -> None:
+        super().__init__(
+            hidden_dim=config.hidden_dim,
+            action_dim=config.action_dim,
+            num_heads=config.num_heads,
+            dropout=config.dropout,
+            bank_size=max(1, config.scene_bank_size),
+            history_steps=max(1, config.scene_history_steps),
+            max_history_steps=config.max_history_steps,
+            write_threshold=config.memory_write_threshold,
+            suppression_margin=config.memory_suppression_margin,
+        )
+
+
+class BeliefMemory(_SelectiveMemoryBank):
+    def __init__(self, config: ObservationMemoryConfig) -> None:
+        super().__init__(
+            hidden_dim=config.hidden_dim,
+            action_dim=config.action_dim,
+            num_heads=config.num_heads,
+            dropout=config.dropout,
+            bank_size=max(1, config.belief_bank_size),
+            history_steps=max(1, config.belief_history_steps),
+            max_history_steps=config.max_history_steps,
+            write_threshold=config.memory_write_threshold + 0.05,
+            suppression_margin=config.memory_suppression_margin,
+        )
+
+
+class DualObservationMemory(nn.Module):
+    def __init__(self, config: ObservationMemoryConfig) -> None:
+        super().__init__()
+        self.scene_memory = SceneMemory(config)
+        self.belief_memory = BeliefMemory(config)
+        self.uncertainty_head = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, 1),
+        )
+
+    def forward(
+        self,
+        scene_tokens: Tensor,
+        history_scene_tokens: Tensor | None = None,
+        history_actions: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        pooled_current = scene_tokens.mean(dim=1)
+        scene_output = self.scene_memory(
+            pooled_current=pooled_current,
+            history_scene_tokens=history_scene_tokens,
+            history_actions=history_actions,
+        )
+        belief_output = self.belief_memory(
+            pooled_current=pooled_current,
+            history_scene_tokens=history_scene_tokens,
+            history_actions=history_actions,
+        )
+        memory_tokens = torch.cat([scene_output["memory_tokens"], belief_output["memory_tokens"]], dim=1)
+        memory_token = memory_tokens.mean(dim=1, keepdim=True)
+        memory_state = torch.cat([scene_output["memory_state"], belief_output["memory_state"]], dim=-1)
+        pooled_memory = memory_tokens.mean(dim=1)
+        return {
+            "scene_memory_tokens": scene_output["memory_tokens"],
+            "belief_memory_tokens": belief_output["memory_tokens"],
+            "memory_tokens": memory_tokens,
+            "memory_token": memory_token,
+            "memory_sequence": torch.cat(
+                [scene_output["memory_sequence"], belief_output["memory_sequence"]],
+                dim=1,
+            ),
+            "memory_state": memory_state,
+            "memory_uncertainty": torch.nn.functional.softplus(self.uncertainty_head(pooled_memory)).squeeze(-1),
+            "memory_write_rate": 0.5 * (scene_output["write_gate"] + belief_output["write_gate"]),
+            "memory_saturation": 0.5 * (scene_output["saturation"] + belief_output["saturation"]),
+            "scene_write_gate": scene_output["write_gate"],
+            "belief_write_gate": belief_output["write_gate"],
+            "memory_scene_state": scene_output["memory_state"],
+            "memory_belief_state": belief_output["memory_state"],
+        }

code/reveal_vla_bimanual/models/planner.py

@@ -24,6 +24,14 @@ class PlannerConfig:
     num_layers: int = 2
     num_phases: int = 5
     num_arm_roles: int = 4
+    top_k: int = 4
+    belief_gain_weight: float = 1.0
+    visibility_gain_weight: float = 0.75
+    clearance_weight: float = 0.75
+    occluder_contact_weight: float = 0.5
+    grasp_affordance_weight: float = 0.75
+    support_stability_weight: float = 0.5
+    residual_weight: float = 0.5
 
 
 class RevealPlanner(nn.Module):
@@ -202,3 +210,186 @@ class InteractionPlanner(nn.Module):
             "best_indices": best_idx,
             "best_chunk": candidate_chunks[batch_indices, best_idx],
         }
+
+
+class StructuredElasticUtility(nn.Module):
+    def __init__(self, config: PlannerConfig) -> None:
+        super().__init__()
+        self.config = config
+
+    def _field_mean(self, tensor: Tensor) -> Tensor:
+        if tensor.ndim == 6:
+            return tensor.mean(dim=(-1, -2, -3))
+        if tensor.ndim == 5:
+            return tensor.mean(dim=(-1, -2))
+        if tensor.ndim == 4:
+            return tensor.mean(dim=(-1, -2))
+        return tensor
+
+    def _initial_scalar(self, state: dict[str, Tensor], key: str) -> Tensor:
+        value = state[key]
+        if value.ndim >= 4:
+            return value.mean(dim=tuple(range(1, value.ndim)))
+        if value.ndim == 3:
+            return value.mean(dim=(-1, -2))
+        if value.ndim == 2:
+            return value.mean(dim=-1)
+        return value
+
+    def forward(
+        self,
+        initial_state: dict[str, Tensor],
+        rollout_state: dict[str, Tensor],
+        candidate_chunks: Tensor,
+    ) -> dict[str, Tensor]:
+        initial_belief = self._initial_scalar(initial_state, "target_belief_field").unsqueeze(1)
+        initial_visibility = self._initial_scalar(initial_state, "visibility_field").unsqueeze(1)
+        belief_future = self._field_mean(rollout_state["target_belief_field"]).mean(dim=-1)
+        visibility_future = self._field_mean(rollout_state["visibility_field"]).mean(dim=-1)
+        clearance = self._field_mean(rollout_state["clearance_field"]).mean(dim=-1)
+        occluder_contact = self._field_mean(rollout_state["occluder_contact_field"]).mean(dim=-1)
+        grasp_affordance = self._field_mean(rollout_state["grasp_affordance_field"]).mean(dim=-1)
+        support_stability = torch.sigmoid(self._field_mean(rollout_state["support_stability_field"])).mean(dim=-1)
+        persistence = self._field_mean(rollout_state["persistence_field"]).mean(dim=-1)
+        reocclusion = self._field_mean(rollout_state["reocclusion_field"]).mean(dim=-1)
+        disturbance = self._field_mean(rollout_state["disturbance_field"]).mean(dim=-1)
+        access_quality = torch.sigmoid(self._field_mean(rollout_state["access_field"])).mean(dim=-1)
+        retrieve_progress = torch.sigmoid(candidate_chunks[:, :, :, -1]).mean(dim=-1)
+        utility = (
+            self.config.belief_gain_weight * (belief_future - initial_belief)
+            + self.config.visibility_gain_weight * (visibility_future - initial_visibility)
+            + self.config.clearance_weight * clearance
+            + self.config.occluder_contact_weight * occluder_contact
+            + self.config.grasp_affordance_weight * grasp_affordance
+            + self.config.persistence_weight * persistence
+            + self.config.support_stability_weight * support_stability
+            + self.config.corridor_weight * access_quality
+            + self.config.task_progress_weight * retrieve_progress
+            - self.config.reocclusion_weight * reocclusion
+            - self.config.disturbance_weight * disturbance
+            - self.config.visibility_weight * (1.0 - visibility_future)
+        )
+        return {
+            "belief_gain": belief_future - initial_belief,
+            "visibility_gain": visibility_future - initial_visibility,
+            "clearance": clearance,
+            "occluder_contact_quality": occluder_contact,
+            "grasp_affordance": grasp_affordance,
+            "persistence": persistence,
+            "support_stability": support_stability,
+            "reocclusion_penalty": reocclusion,
+            "disturbance_penalty": disturbance,
+            "access_quality": access_quality,
+            "task_progress": retrieve_progress,
+            "utility_structured": utility,
+        }
+
+
+class ResidualPlannerScorer(nn.Module):
+    def __init__(self, config: PlannerConfig) -> None:
+        super().__init__()
+        feature_dim = (config.action_dim * 2) + 11
+        self.trunk = nn.Sequential(
+            nn.LayerNorm(feature_dim),
+            nn.Linear(feature_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.success_head = nn.Linear(config.hidden_dim, 1)
+        self.risk_head = nn.Linear(config.hidden_dim, 1)
+        self.residual_head = nn.Linear(config.hidden_dim, 1)
+
+    def forward(
+        self,
+        candidate_chunks: Tensor,
+        structured: dict[str, Tensor],
+        proposal_logits: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        candidate_mean = candidate_chunks.mean(dim=2)
+        candidate_terminal = candidate_chunks[:, :, -1]
+        components = torch.stack(
+            [
+                structured["belief_gain"],
+                structured["visibility_gain"],
+                structured["clearance"],
+                structured["occluder_contact_quality"],
+                structured["grasp_affordance"],
+                structured["persistence"],
+                structured["support_stability"],
+                structured["reocclusion_penalty"],
+                structured["disturbance_penalty"],
+                structured["access_quality"],
+                structured["task_progress"],
+            ],
+            dim=-1,
+        )
+        features = torch.cat([candidate_mean, candidate_terminal, components], dim=-1)
+        hidden = self.trunk(features)
+        success_logits = self.success_head(hidden).squeeze(-1)
+        risk_values = torch.sigmoid(self.risk_head(hidden)).squeeze(-1)
+        residual = self.residual_head(hidden).squeeze(-1)
+        if proposal_logits is not None and proposal_logits.shape == residual.shape:
+            residual = residual + 0.25 * proposal_logits.sigmoid()
+        return {
+            "planner_hidden": hidden,
+            "success_logits": success_logits,
+            "risk_values": risk_values,
+            "utility_residual": residual,
+        }
+
+
+class CascadePlanner(nn.Module):
+    def __init__(self, config: PlannerConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.structured = StructuredElasticUtility(config)
+        self.residual = ResidualPlannerScorer(config)
+
+    def shortlist(self, proposal_logits: Tensor | None, candidate_chunks: Tensor) -> Tensor:
+        batch_size, num_candidates = candidate_chunks.shape[:2]
+        top_k = min(max(1, self.config.top_k), num_candidates)
+        if proposal_logits is None:
+            cheap_scores = -candidate_chunks.square().mean(dim=(-1, -2))
+        else:
+            cheap_scores = proposal_logits
+        return cheap_scores.topk(top_k, dim=-1).indices
+
+    def select_best(
+        self,
+        initial_state: dict[str, Tensor],
+        candidate_chunks: Tensor,
+        rollout_state: dict[str, Tensor],
+        proposal_logits: Tensor | None = None,
+        candidate_indices: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        structured = self.structured(
+            initial_state=initial_state,
+            rollout_state=rollout_state,
+            candidate_chunks=candidate_chunks,
+        )
+        residual = self.residual(
+            candidate_chunks=candidate_chunks,
+            structured=structured,
+            proposal_logits=proposal_logits,
+        )
+        utility_total = structured["utility_structured"] + self.config.residual_weight * residual["utility_residual"]
+        utility_total = utility_total + residual["success_logits"].sigmoid() - residual["risk_values"]
+        best_local = utility_total.argmax(dim=-1)
+        batch_indices = torch.arange(candidate_chunks.shape[0], device=candidate_chunks.device)
+        if candidate_indices is None:
+            best_indices = best_local
+        else:
+            best_indices = candidate_indices[batch_indices, best_local]
+        return {
+            **structured,
+            **residual,
+            "utility_total": utility_total,
+            "utility_scores": utility_total,
+            "best_indices": best_indices,
+            "best_chunk": candidate_chunks[batch_indices, best_local],
+            "ranking_diagnostics": {
+                "topk_indices": candidate_indices if candidate_indices is not None else best_local.unsqueeze(-1),
+                "best_local_indices": best_local,
+            },
+        }

code/reveal_vla_bimanual/models/policy.py

@@ -6,13 +6,28 @@ from typing import Sequence
 import torch
 from torch import Tensor, nn
 
-from models.action_decoder import ACTBimanualChunkDecoder, ChunkDecoderConfig, InteractionChunkDecoder
+from models.action_decoder import (
+    ACTBimanualChunkDecoder,
+    ChunkDecoderConfig,
+    InteractionChunkDecoder,
+    SymmetricCoordinatedChunkDecoder,
+)
 from models.backbones import FrozenVLBackbone, FrozenVLBackboneConfig
 from models.multiview_fusion import MultiViewFusion, MultiViewFusionConfig
-from models.observation_memory import InteractionObservationMemory, ObservationMemory, ObservationMemoryConfig
-from models.planner import InteractionPlanner, PlannerConfig, RevealPlanner
-from models.reveal_head import InteractionStateHead, RevealHeadConfig, RevealStateHead
-from models.world_model import InteractionRolloutModel, RevealWM, RevealWMConfig
+from models.observation_memory import (
+    DualObservationMemory,
+    InteractionObservationMemory,
+    ObservationMemory,
+    ObservationMemoryConfig,
+)
+from models.planner import CascadePlanner, InteractionPlanner, PlannerConfig, RevealPlanner
+from models.reveal_head import (
+    ElasticOcclusionStateHead,
+    InteractionStateHead,
+    RevealHeadConfig,
+    RevealStateHead,
+)
+from models.world_model import ElasticOcclusionWorldModel, InteractionRolloutModel, RevealWM, RevealWMConfig
 
 
 @dataclass
@@ -351,3 +366,302 @@ class InteractionBimanualPolicy(BackboneOnlyPolicy):
             outputs["planner_scores"] = selected["utility_scores"]
             outputs["best_candidate_indices"] = selected["best_indices"]
         return outputs
+
+
+class ElasticRevealBimanualPolicy(BackboneOnlyPolicy):
+    def __init__(self, config: PolicyConfig) -> None:
+        super().__init__(config)
+        self.memory = DualObservationMemory(config.memory)
+        self.decoder = SymmetricCoordinatedChunkDecoder(config.decoder)
+        self.elastic_state_head = ElasticOcclusionStateHead(config.reveal_head)
+        self.world_model = ElasticOcclusionWorldModel(config.world_model)
+        self.planner = CascadePlanner(config.planner)
+
+    def _encode_scene_with_optional_depth(
+        self,
+        images: Tensor,
+        proprio: Tensor,
+        texts: Sequence[str] | None = None,
+        language_tokens: dict[str, Tensor] | None = None,
+        depths: Tensor | None = None,
+        depth_valid: Tensor | None = None,
+        camera_intrinsics: Tensor | None = None,
+        camera_extrinsics: Tensor | None = None,
+        use_depth: bool = True,
+    ) -> dict[str, Tensor]:
+        encoded = self.backbone.encode_images(
+            images,
+            depths=depths if use_depth else None,
+            depth_valid=depth_valid if use_depth else None,
+            camera_intrinsics=camera_intrinsics if use_depth else None,
+            camera_extrinsics=camera_extrinsics if use_depth else None,
+            return_aux=True,
+        )
+        assert isinstance(encoded, dict)
+        text_tokens = self._encode_language(images, texts=texts, language_tokens=language_tokens)
+        fused = self.fusion(
+            image_tokens=encoded["rgb_tokens"],
+            proprio=proprio,
+            language_tokens=text_tokens,
+            depth_tokens=encoded.get("depth_tokens"),
+            camera_tokens=encoded.get("camera_tokens"),
+            return_aux=True,
+        )
+        assert isinstance(fused, dict)
+        return {
+            "scene_tokens": fused["scene_tokens"],
+            "view_summaries": fused["view_summaries"],
+            "geometry_summaries": fused["geometry_summaries"],
+            "depth_tokens": encoded.get("depth_tokens"),
+            "camera_tokens": encoded.get("camera_tokens"),
+        }
+
+    def _expand_language_tokens_for_history(
+        self,
+        language_tokens: dict[str, Tensor] | None,
+        history_steps: int,
+    ) -> dict[str, Tensor] | None:
+        if language_tokens is None:
+            return None
+        return {
+            key: value.unsqueeze(1).expand(-1, history_steps, *value.shape[1:]).reshape(
+                value.shape[0] * history_steps, *value.shape[1:]
+            )
+            for key, value in language_tokens.items()
+        }
+
+    def encode_history_with_optional_depth(
+        self,
+        history_images: Tensor | None,
+        history_proprio: Tensor | None,
+        texts: Sequence[str] | None = None,
+        language_tokens: dict[str, Tensor] | None = None,
+        history_depths: Tensor | None = None,
+        history_depth_valid: Tensor | None = None,
+        camera_intrinsics: Tensor | None = None,
+        camera_extrinsics: Tensor | None = None,
+        use_depth: bool = True,
+    ) -> Tensor | None:
+        if history_images is None or history_proprio is None or history_images.numel() == 0:
+            return None
+        batch_size, history_steps = history_images.shape[:2]
+        flat_images = history_images.reshape(batch_size * history_steps, *history_images.shape[2:])
+        flat_proprio = history_proprio.reshape(batch_size * history_steps, history_proprio.shape[-1])
+        flat_depths = None
+        flat_depth_valid = None
+        if history_depths is not None and history_depths.numel() > 0:
+            flat_depths = history_depths.reshape(batch_size * history_steps, *history_depths.shape[2:])
+        if history_depth_valid is not None and history_depth_valid.numel() > 0:
+            flat_depth_valid = history_depth_valid.reshape(batch_size * history_steps, *history_depth_valid.shape[2:])
+        if language_tokens is None:
+            flat_texts = [text for text in texts for _ in range(history_steps)] if texts is not None else None
+            flat_language_tokens = None
+        else:
+            flat_texts = None
+            flat_language_tokens = self._expand_language_tokens_for_history(language_tokens, history_steps)
+        history_scene = self._encode_scene_with_optional_depth(
+            images=flat_images,
+            proprio=flat_proprio,
+            texts=flat_texts,
+            language_tokens=flat_language_tokens,
+            depths=flat_depths,
+            depth_valid=flat_depth_valid,
+            camera_intrinsics=None,
+            camera_extrinsics=None,
+            use_depth=use_depth,
+        )["scene_tokens"]
+        return history_scene.view(batch_size, history_steps, history_scene.shape[1], history_scene.shape[2])
+
+    def _tile_tensor(self, value: Tensor, num_candidates: int) -> Tensor:
+        return value.unsqueeze(1).expand(-1, num_candidates, *value.shape[1:]).reshape(
+            value.shape[0] * num_candidates,
+            *value.shape[1:],
+        )
+
+    def _tile_state(self, state: dict[str, Tensor], num_candidates: int) -> dict[str, Tensor]:
+        tiled: dict[str, Tensor] = {}
+        for key, value in state.items():
+            if isinstance(value, Tensor):
+                tiled[key] = self._tile_tensor(value, num_candidates)
+        return tiled
+
+    def forward(
+        self,
+        images: Tensor,
+        proprio: Tensor,
+        texts: Sequence[str] | None = None,
+        language_tokens: dict[str, Tensor] | None = None,
+        history_images: Tensor | None = None,
+        history_proprio: Tensor | None = None,
+        history_actions: Tensor | None = None,
+        plan: bool = True,
+        support_mode_conditioning: bool = True,
+        candidate_chunks_override: Tensor | None = None,
+        use_depth: bool = True,
+        use_world_model: bool = True,
+        use_planner: bool = True,
+        use_role_tokens: bool = True,
+        history_steps_override: int | None = None,
+        depths: Tensor | None = None,
+        depth_valid: Tensor | None = None,
+        camera_intrinsics: Tensor | None = None,
+        camera_extrinsics: Tensor | None = None,
+        history_depths: Tensor | None = None,
+        history_depth_valid: Tensor | None = None,
+        compute_equivariance_probe: bool = False,
+    ) -> dict[str, Tensor]:
+        scene_output = self._encode_scene_with_optional_depth(
+            images=images,
+            proprio=proprio,
+            texts=texts,
+            language_tokens=language_tokens,
+            depths=depths,
+            depth_valid=depth_valid,
+            camera_intrinsics=camera_intrinsics,
+            camera_extrinsics=camera_extrinsics,
+            use_depth=use_depth,
+        )
+        scene_tokens = scene_output["scene_tokens"]
+        history_scene_tokens = self.encode_history_with_optional_depth(
+            history_images=history_images,
+            history_proprio=history_proprio,
+            texts=texts,
+            language_tokens=language_tokens,
+            history_depths=history_depths,
+            history_depth_valid=history_depth_valid,
+            camera_intrinsics=camera_intrinsics,
+            camera_extrinsics=camera_extrinsics,
+            use_depth=use_depth,
+        )
+        if history_steps_override is not None and history_scene_tokens is not None and history_scene_tokens.numel() > 0:
+            history_scene_tokens = history_scene_tokens[:, -history_steps_override:]
+            if history_actions is not None and history_actions.numel() > 0:
+                history_actions = history_actions[:, -history_steps_override:]
+        memory_output = self.memory(
+            scene_tokens,
+            history_scene_tokens=history_scene_tokens,
+            history_actions=history_actions,
+        )
+        elastic_state = self.elastic_state_head(
+            scene_tokens,
+            memory_tokens=memory_output["memory_tokens"],
+        )
+        elastic_state["memory_tokens"] = memory_output["memory_tokens"]
+        elastic_state["memory_token"] = memory_output["memory_token"]
+        elastic_state["scene_memory_tokens"] = memory_output["scene_memory_tokens"]
+        elastic_state["belief_memory_tokens"] = memory_output["belief_memory_tokens"]
+        if not use_role_tokens:
+            elastic_state = dict(elastic_state)
+            elastic_state["arm_role_logits"] = torch.zeros_like(elastic_state["arm_role_logits"])
+
+        decoded = self.decoder(
+            scene_tokens,
+            interaction_state=elastic_state,
+            memory_tokens=memory_output["memory_tokens"],
+            compute_equivariance_probe=compute_equivariance_probe,
+        )
+        outputs = {
+            **decoded,
+            "scene_tokens": scene_tokens,
+            "history_scene_tokens": history_scene_tokens,
+            "memory_output": memory_output,
+            "memory_uncertainty": memory_output["memory_uncertainty"],
+            "interaction_state": elastic_state,
+            "reveal_state": elastic_state,
+            "view_summaries": scene_output["view_summaries"],
+            "geometry_summaries": scene_output["geometry_summaries"],
+        }
+
+        candidate_chunks = candidate_chunks_override
+        proposal_logits = outputs.get("proposal_logits")
+        if candidate_chunks is None:
+            candidate_chunks = self.decoder.sample_candidates(
+                outputs["action_mean"],
+                outputs["action_log_std"],
+                num_candidates=self.config.decoder.num_candidates,
+                proposal_candidates=outputs.get("proposal_candidates"),
+            )
+        else:
+            proposal_logits = None
+        outputs["candidate_chunks"] = candidate_chunks
+
+        if not plan or not use_planner:
+            outputs["planned_chunk"] = outputs["action_mean"]
+            outputs["planned_rollout"] = {}
+            outputs["planner_success_logits"] = torch.zeros(
+                candidate_chunks.shape[:2],
+                device=candidate_chunks.device,
+                dtype=candidate_chunks.dtype,
+            )
+            outputs["planner_risk_values"] = torch.zeros_like(outputs["planner_success_logits"])
+            outputs["planner_scores"] = torch.zeros_like(outputs["planner_success_logits"])
+            outputs["best_candidate_indices"] = torch.zeros(
+                candidate_chunks.shape[0],
+                dtype=torch.long,
+                device=candidate_chunks.device,
+            )
+            return outputs
+
+        shortlist_indices = self.planner.shortlist(proposal_logits=proposal_logits, candidate_chunks=candidate_chunks)
+        outputs["planner_topk_indices"] = shortlist_indices
+        batch_size = candidate_chunks.shape[0]
+        batch_indices = torch.arange(batch_size, device=candidate_chunks.device).unsqueeze(-1)
+        topk_candidates = candidate_chunks[batch_indices, shortlist_indices]
+        outputs["planner_topk_candidates"] = topk_candidates
+        if proposal_logits is not None:
+            topk_proposal_logits = proposal_logits.gather(1, shortlist_indices)
+        else:
+            topk_proposal_logits = None
+
+        if not use_world_model:
+            score_source = topk_proposal_logits if topk_proposal_logits is not None else -topk_candidates.square().mean(dim=(-1, -2))
+            best_local = score_source.argmax(dim=-1)
+            best_indices = shortlist_indices[torch.arange(batch_size, device=best_local.device), best_local]
+            outputs["planned_chunk"] = candidate_chunks[torch.arange(batch_size, device=best_local.device), best_indices]
+            outputs["planned_rollout"] = {}
+            outputs["planner_success_logits"] = torch.zeros_like(score_source)
+            outputs["planner_risk_values"] = torch.zeros_like(score_source)
+            outputs["planner_scores"] = score_source
+            outputs["best_candidate_indices"] = best_indices
+            outputs["utility_structured"] = score_source
+            outputs["utility_residual"] = torch.zeros_like(score_source)
+            outputs["utility_total"] = score_source
+            return outputs
+
+        num_topk = topk_candidates.shape[1]
+        flat_chunks = topk_candidates.view(batch_size * num_topk, topk_candidates.shape[2], topk_candidates.shape[3])
+        tiled_scene = self._tile_tensor(scene_tokens, num_topk)
+        planning_state = elastic_state
+        if not support_mode_conditioning:
+            planning_state = dict(elastic_state)
+            planning_state["support_mode_logits"] = torch.zeros_like(elastic_state["support_mode_logits"])
+        tiled_state = self._tile_state(planning_state, num_topk)
+        rollout = self.world_model(
+            scene_tokens=tiled_scene,
+            interaction_state=tiled_state,
+            action_chunk=flat_chunks,
+            memory_tokens=self._tile_tensor(memory_output["memory_tokens"], num_topk),
+            scene_memory_tokens=self._tile_tensor(memory_output["scene_memory_tokens"], num_topk),
+            belief_memory_tokens=self._tile_tensor(memory_output["belief_memory_tokens"], num_topk),
+        )
+        reshaped_rollout = {
+            key: value.view(batch_size, num_topk, *value.shape[1:]) for key, value in rollout.items()
+        }
+        selected = self.planner.select_best(
+            initial_state=elastic_state,
+            candidate_chunks=topk_candidates,
+            rollout_state=reshaped_rollout,
+            proposal_logits=topk_proposal_logits,
+            candidate_indices=shortlist_indices,
+        )
+        outputs["planned_rollout"] = reshaped_rollout
+        outputs["planned_chunk"] = selected["best_chunk"]
+        outputs["planner_success_logits"] = selected["success_logits"]
+        outputs["planner_risk_values"] = selected["risk_values"]
+        outputs["planner_scores"] = selected["utility_total"]
+        outputs["best_candidate_indices"] = selected["best_indices"]
+        outputs["utility_structured"] = selected["utility_structured"]
+        outputs["utility_residual"] = selected["utility_residual"]
+        outputs["utility_total"] = selected["utility_total"]
+        outputs["ranking_diagnostics"] = selected["ranking_diagnostics"]
+        return outputs

code/reveal_vla_bimanual/models/reveal_head.py

@@ -317,3 +317,245 @@ class InteractionStateHead(nn.Module):
             scene_tokens=scene_tokens,
             memory_tokens=memory_tokens,
         )
+
+
+class ElasticOcclusionFieldDecoder(nn.Module):
+    def __init__(self, config: RevealHeadConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.field_queries = nn.Parameter(
+            torch.randn(config.field_size * config.field_size, config.hidden_dim) * 0.02
+        )
+        self.field_attention = nn.MultiheadAttention(
+            embed_dim=config.hidden_dim,
+            num_heads=config.num_heads,
+            batch_first=True,
+        )
+        self.field_mlp = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+        )
+        summary_dim = config.hidden_dim * 4
+        self.summary_proj = nn.Sequential(
+            nn.LayerNorm(summary_dim),
+            nn.Linear(summary_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.phase_head = nn.Sequential(
+            nn.LayerNorm(summary_dim),
+            nn.Linear(summary_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.num_phases),
+        )
+        self.arm_role_head = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim * 2),
+            nn.Linear(config.hidden_dim * 2, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.num_arm_roles),
+        )
+        self.arm_identity = nn.Embedding(2, config.hidden_dim)
+        self.support_mode = nn.Sequential(
+            nn.LayerNorm(summary_dim),
+            nn.Linear(summary_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.num_support_modes),
+        )
+        self.access_field = nn.Conv2d(config.hidden_dim, config.num_support_modes, kernel_size=1)
+        self.target_belief_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.visibility_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.clearance_field = nn.Conv2d(config.hidden_dim, 2, kernel_size=1)
+        self.occluder_contact_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.grasp_affordance_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.support_stability_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.persistence_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.reocclusion_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.disturbance_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.uncertainty_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.reocclusion_head = nn.Sequential(
+            nn.LayerNorm(summary_dim),
+            nn.Linear(summary_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.num_support_modes),
+        )
+
+    def _pool_source(self, source_tokens: Tensor | None, fallback: Tensor) -> Tensor:
+        if source_tokens is None or source_tokens.numel() == 0:
+            return fallback.new_zeros(fallback.shape)
+        return source_tokens.mean(dim=1)
+
+    def _field_mean(self, field: Tensor) -> Tensor:
+        return field.mean(dim=(-1, -2))
+
+    def _upsampled_belief(self, target_belief_field: Tensor) -> Tensor:
+        if target_belief_field.shape[-1] == self.config.belief_map_size:
+            return target_belief_field
+        return F.interpolate(
+            target_belief_field,
+            size=(self.config.belief_map_size, self.config.belief_map_size),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+    def forward(
+        self,
+        interaction_tokens: Tensor,
+        scene_tokens: Tensor | None = None,
+        memory_tokens: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        batch_size = interaction_tokens.shape[0]
+        pooled_interaction = interaction_tokens.mean(dim=1)
+        pooled_scene = self._pool_source(scene_tokens, pooled_interaction)
+        pooled_memory = self._pool_source(memory_tokens, pooled_interaction)
+
+        field_queries = self.field_queries.unsqueeze(0).expand(batch_size, -1, -1)
+        source_tokens = interaction_tokens
+        if scene_tokens is not None:
+            source_tokens = torch.cat([source_tokens, scene_tokens], dim=1)
+        if memory_tokens is not None:
+            source_tokens = torch.cat([source_tokens, memory_tokens], dim=1)
+        field_tokens, _ = self.field_attention(field_queries, source_tokens, source_tokens)
+        field_tokens = field_tokens + self.field_mlp(field_tokens)
+
+        side = self.config.field_size
+        grid = field_tokens.transpose(1, 2).reshape(batch_size, self.config.hidden_dim, side, side)
+        pooled_field = field_tokens.mean(dim=1)
+        summary_input = torch.cat([pooled_interaction, pooled_field, pooled_scene, pooled_memory], dim=-1)
+        latent_summary = self.summary_proj(summary_input)
+
+        access_field = self.access_field(grid)
+        target_belief_field = self.target_belief_field(grid)
+        visibility_field = self.visibility_field(grid)
+        clearance_field = self.clearance_field(grid)
+        occluder_contact_field = self.occluder_contact_field(grid)
+        grasp_affordance_field = self.grasp_affordance_field(grid)
+        support_stability_field = self.support_stability_field(grid)
+        persistence_field = torch.sigmoid(self.persistence_field(grid))
+        reocclusion_field = torch.sigmoid(self.reocclusion_field(grid))
+        disturbance_field = torch.sigmoid(self.disturbance_field(grid))
+        uncertainty_field = F.softplus(self.uncertainty_field(grid))
+
+        support_stability_prob = torch.sigmoid(support_stability_field)
+        risk_field = torch.sigmoid(
+            disturbance_field
+            + 0.75 * reocclusion_field
+            + 0.5 * (1.0 - support_stability_prob)
+            + 0.25 * uncertainty_field
+        )
+        corridor_source = access_field.amax(dim=-2)
+        corridor_logits = F.interpolate(
+            corridor_source,
+            size=self.config.num_approach_templates,
+            mode="linear",
+            align_corners=False,
+        )
+        access_prob = torch.sigmoid(access_field)
+        weighted_persistence = (persistence_field.expand_as(access_prob) * access_prob).sum(dim=(-1, -2))
+        access_mass = access_prob.sum(dim=(-1, -2)).clamp_min(1e-4)
+        persistence_horizon = self.config.rollout_horizon * weighted_persistence / access_mass
+        disturbance_cost = disturbance_field.mean(dim=(-1, -2)).squeeze(1)
+
+        arm_identity = self.arm_identity.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        arm_tokens = pooled_interaction.unsqueeze(1).expand(-1, 2, -1) + arm_identity
+        arm_role_input = torch.cat(
+            [arm_tokens, latent_summary.unsqueeze(1).expand(-1, 2, -1)],
+            dim=-1,
+        )
+        arm_role_logits = self.arm_role_head(arm_role_input)
+        target_belief_map = self._upsampled_belief(target_belief_field)
+        compact_components = [
+            target_belief_field.mean(dim=(-1, -2)).squeeze(1),
+            visibility_field.mean(dim=(-1, -2)).squeeze(1),
+            clearance_field.mean(dim=(-1, -2)).mean(dim=1),
+            occluder_contact_field.mean(dim=(-1, -2)).squeeze(1),
+            grasp_affordance_field.mean(dim=(-1, -2)).squeeze(1),
+            support_stability_prob.mean(dim=(-1, -2)).squeeze(1),
+            persistence_field.mean(dim=(-1, -2)).squeeze(1),
+            reocclusion_field.mean(dim=(-1, -2)).squeeze(1),
+            disturbance_field.mean(dim=(-1, -2)).squeeze(1),
+            risk_field.mean(dim=(-1, -2)).squeeze(1),
+            uncertainty_field.mean(dim=(-1, -2)).squeeze(1),
+            access_prob.mean(dim=(-1, -2)).transpose(0, 1).transpose(0, 1),
+            self.support_mode(summary_input),
+            self.phase_head(summary_input),
+            arm_role_logits.reshape(batch_size, -1),
+        ]
+        compact_state = torch.cat(
+            [component if component.ndim > 1 else component.unsqueeze(-1) for component in compact_components],
+            dim=-1,
+        )
+
+        output = {
+            "phase_logits": self.phase_head(summary_input),
+            "arm_role_logits": arm_role_logits,
+            "target_belief_field": target_belief_field,
+            "visibility_field": visibility_field,
+            "clearance_field": clearance_field,
+            "occluder_contact_field": occluder_contact_field,
+            "grasp_affordance_field": grasp_affordance_field,
+            "support_stability_field": support_stability_field,
+            "persistence_field": persistence_field,
+            "reocclusion_field": reocclusion_field,
+            "disturbance_field": disturbance_field,
+            "risk_field": risk_field,
+            "uncertainty_field": uncertainty_field,
+            "interaction_tokens": interaction_tokens,
+            "field_tokens": field_tokens,
+            "latent_summary": latent_summary,
+            "support_mode_logits": self.support_mode(summary_input),
+            "corridor_logits": corridor_logits,
+            "persistence_horizon": persistence_horizon,
+            "disturbance_cost": disturbance_cost,
+            "belief_map": target_belief_map,
+            "reocclusion_logit": self.reocclusion_head(summary_input),
+            "persistence_uncertainty": uncertainty_field.mean(dim=(-1, -2)).squeeze(1),
+            "access_field": access_field,
+            "uncertainty": uncertainty_field.mean(dim=(-1, -2)).squeeze(1),
+            "compact_state": compact_state,
+        }
+        output["target_field"] = output["target_belief_field"]
+        output["actor_feasibility_field"] = output["clearance_field"]
+        return output
+
+
+class ElasticOcclusionStateHead(nn.Module):
+    def __init__(self, config: RevealHeadConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.interaction_queries = nn.Parameter(
+            torch.randn(config.num_interaction_tokens, config.hidden_dim) * 0.02
+        )
+        self.interaction_attention = nn.MultiheadAttention(
+            embed_dim=config.hidden_dim,
+            num_heads=config.num_heads,
+            batch_first=True,
+        )
+        self.interaction_mlp = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+        )
+        self.decoder = ElasticOcclusionFieldDecoder(config)
+
+    def forward(
+        self,
+        scene_tokens: Tensor,
+        memory_token: Tensor | None = None,
+        memory_tokens: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        if memory_tokens is None:
+            memory_tokens = memory_token
+        source_tokens = scene_tokens
+        if memory_tokens is not None:
+            source_tokens = torch.cat([source_tokens, memory_tokens], dim=1)
+        batch_size = source_tokens.shape[0]
+        interaction_queries = self.interaction_queries.unsqueeze(0).expand(batch_size, -1, -1)
+        interaction_tokens, _ = self.interaction_attention(interaction_queries, source_tokens, source_tokens)
+        interaction_tokens = interaction_tokens + self.interaction_mlp(interaction_tokens)
+        return self.decoder(
+            interaction_tokens=interaction_tokens,
+            scene_tokens=scene_tokens,
+            memory_tokens=memory_tokens,
+        )

code/reveal_vla_bimanual/models/world_model.py

@@ -22,6 +22,8 @@ class RevealWMConfig:
     num_interaction_tokens: int = 8
     belief_map_size: int = 32
     predict_belief_map: bool = True
+    scene_bank_size: int = 2
+    belief_bank_size: int = 2
 
 
 class RevealWM(nn.Module):
@@ -152,3 +154,186 @@ class InteractionRolloutModel(nn.Module):
         for key, values in outputs.items():
             stacked[key] = torch.stack(values, dim=1)
         return stacked
+
+
+class ElasticOcclusionWorldModel(nn.Module):
+    def __init__(self, config: RevealWMConfig) -> None:
+        super().__init__()
+        self.config = config
+        compact_state_dim = (
+            11
+            + config.num_support_modes
+            + config.num_support_modes
+            + config.num_phases
+            + (2 * config.num_arm_roles)
+        )
+        self.state_encoder = nn.Sequential(
+            nn.LayerNorm(compact_state_dim),
+            nn.Linear(compact_state_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.scene_memory_proj = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.belief_memory_proj = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.action_encoder = nn.Sequential(
+            nn.LayerNorm(config.action_dim),
+            nn.Linear(config.action_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.transition = nn.GRUCell(config.hidden_dim * 4, config.hidden_dim)
+        self.scene_memory_update = nn.Linear(config.hidden_dim, config.hidden_dim)
+        self.belief_memory_update = nn.Linear(config.hidden_dim, config.hidden_dim)
+        self.compact_decoder = nn.Linear(config.hidden_dim, compact_state_dim)
+        field_elements = config.field_size * config.field_size
+        self.target_belief_head = nn.Linear(config.hidden_dim, field_elements)
+        self.visibility_head = nn.Linear(config.hidden_dim, field_elements)
+        self.clearance_head = nn.Linear(config.hidden_dim, 2 * field_elements)
+        self.occluder_contact_head = nn.Linear(config.hidden_dim, field_elements)
+        self.grasp_affordance_head = nn.Linear(config.hidden_dim, field_elements)
+        self.support_stability_head = nn.Linear(config.hidden_dim, field_elements)
+        self.persistence_head = nn.Linear(config.hidden_dim, field_elements)
+        self.reocclusion_head = nn.Linear(config.hidden_dim, field_elements)
+        self.disturbance_head = nn.Linear(config.hidden_dim, field_elements)
+        self.uncertainty_head = nn.Linear(config.hidden_dim, field_elements)
+        self.access_head = nn.Linear(config.hidden_dim, config.num_support_modes * field_elements)
+
+    def _compact_from_state(self, interaction_state: dict[str, Tensor]) -> Tensor:
+        if "compact_state" in interaction_state:
+            return interaction_state["compact_state"]
+        components = [
+            interaction_state["target_belief_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["visibility_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["clearance_field"].mean(dim=(-1, -2)).mean(dim=1),
+            interaction_state["occluder_contact_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["grasp_affordance_field"].mean(dim=(-1, -2)).squeeze(1),
+            torch.sigmoid(interaction_state["support_stability_field"]).mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["persistence_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["reocclusion_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["disturbance_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["risk_field"].mean(dim=(-1, -2)).squeeze(1),
+            interaction_state["uncertainty_field"].mean(dim=(-1, -2)).squeeze(1),
+            torch.sigmoid(interaction_state["access_field"]).mean(dim=(-1, -2)),
+            interaction_state["support_mode_logits"],
+            interaction_state["phase_logits"],
+            interaction_state["arm_role_logits"].reshape(interaction_state["arm_role_logits"].shape[0], -1),
+        ]
+        return torch.cat([component if component.ndim > 1 else component.unsqueeze(-1) for component in components], dim=-1)
+
+    def _decode_fields(self, latent: Tensor) -> dict[str, Tensor]:
+        batch_size = latent.shape[0]
+        side = self.config.field_size
+        target_belief_field = self.target_belief_head(latent).view(batch_size, 1, side, side)
+        visibility_field = self.visibility_head(latent).view(batch_size, 1, side, side)
+        clearance_field = self.clearance_head(latent).view(batch_size, 2, side, side)
+        occluder_contact_field = self.occluder_contact_head(latent).view(batch_size, 1, side, side)
+        grasp_affordance_field = self.grasp_affordance_head(latent).view(batch_size, 1, side, side)
+        support_stability_field = self.support_stability_head(latent).view(batch_size, 1, side, side)
+        persistence_field = torch.sigmoid(self.persistence_head(latent).view(batch_size, 1, side, side))
+        reocclusion_field = torch.sigmoid(self.reocclusion_head(latent).view(batch_size, 1, side, side))
+        disturbance_field = torch.sigmoid(self.disturbance_head(latent).view(batch_size, 1, side, side))
+        uncertainty_field = torch.nn.functional.softplus(self.uncertainty_head(latent).view(batch_size, 1, side, side))
+        risk_field = torch.sigmoid(
+            disturbance_field
+            + 0.75 * reocclusion_field
+            + 0.5 * (1.0 - torch.sigmoid(support_stability_field))
+            + 0.25 * uncertainty_field
+        )
+        access_field = self.access_head(latent).view(batch_size, self.config.num_support_modes, side, side)
+        corridor_source = access_field.amax(dim=-2)
+        corridor_logits = torch.nn.functional.interpolate(
+            corridor_source,
+            size=self.config.num_approach_templates,
+            mode="linear",
+            align_corners=False,
+        )
+        access_prob = torch.sigmoid(access_field)
+        weighted_persistence = (persistence_field.expand_as(access_prob) * access_prob).sum(dim=(-1, -2))
+        access_mass = access_prob.sum(dim=(-1, -2)).clamp_min(1e-4)
+        persistence_horizon = self.config.rollout_horizon * weighted_persistence / access_mass
+        return {
+            "target_belief_field": target_belief_field,
+            "visibility_field": visibility_field,
+            "clearance_field": clearance_field,
+            "occluder_contact_field": occluder_contact_field,
+            "grasp_affordance_field": grasp_affordance_field,
+            "support_stability_field": support_stability_field,
+            "persistence_field": persistence_field,
+            "reocclusion_field": reocclusion_field,
+            "disturbance_field": disturbance_field,
+            "risk_field": risk_field,
+            "uncertainty_field": uncertainty_field,
+            "access_field": access_field,
+            "corridor_logits": corridor_logits,
+            "persistence_horizon": persistence_horizon,
+            "disturbance_cost": disturbance_field.mean(dim=(-1, -2)).squeeze(1),
+            "belief_map": torch.nn.functional.interpolate(
+                target_belief_field,
+                size=(self.config.belief_map_size, self.config.belief_map_size),
+                mode="bilinear",
+                align_corners=False,
+            ),
+            "target_field": target_belief_field,
+            "actor_feasibility_field": clearance_field,
+        }
+
+    def forward(
+        self,
+        scene_tokens: Tensor,
+        interaction_state: dict[str, Tensor],
+        action_chunk: Tensor,
+        memory_tokens: Tensor | None = None,
+        scene_memory_tokens: Tensor | None = None,
+        belief_memory_tokens: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        if scene_memory_tokens is None:
+            scene_memory_tokens = interaction_state.get("scene_memory_tokens")
+        if belief_memory_tokens is None:
+            belief_memory_tokens = interaction_state.get("belief_memory_tokens")
+        if scene_memory_tokens is None and memory_tokens is not None:
+            scene_memory_tokens = memory_tokens
+        if belief_memory_tokens is None and memory_tokens is not None:
+            belief_memory_tokens = memory_tokens
+        if scene_memory_tokens is None:
+            scene_memory_tokens = scene_tokens[:, :1]
+        if belief_memory_tokens is None:
+            belief_memory_tokens = scene_tokens[:, :1]
+
+        latent = self.state_encoder(self._compact_from_state(interaction_state))
+        scene_memory = self.scene_memory_proj(scene_memory_tokens.mean(dim=1))
+        belief_memory = self.belief_memory_proj(belief_memory_tokens.mean(dim=1))
+        outputs: dict[str, list[Tensor]] = {}
+        scene_bias = scene_tokens.mean(dim=1)
+
+        for step in range(action_chunk.shape[1]):
+            action_latent = self.action_encoder(action_chunk[:, step])
+            transition_input = torch.cat([latent, action_latent, scene_memory, belief_memory], dim=-1)
+            latent = self.transition(transition_input, latent + 0.1 * scene_bias)
+            scene_memory = 0.75 * scene_memory + 0.25 * torch.tanh(self.scene_memory_update(latent))
+            belief_memory = 0.65 * belief_memory + 0.35 * torch.tanh(self.belief_memory_update(latent))
+            compact_state = self.compact_decoder(latent)
+            decoded = self._decode_fields(latent)
+            decoded["compact_state"] = compact_state
+            decoded["phase_logits"] = compact_state[:, -(self.config.num_phases + (2 * self.config.num_arm_roles)) : -(2 * self.config.num_arm_roles)]
+            role_slice = compact_state[:, -(2 * self.config.num_arm_roles) :]
+            decoded["arm_role_logits"] = role_slice.view(role_slice.shape[0], 2, self.config.num_arm_roles)
+            decoded["support_mode_logits"] = compact_state[
+                :,
+                -(self.config.num_phases + (2 * self.config.num_arm_roles) + self.config.num_support_modes) : -(self.config.num_phases + (2 * self.config.num_arm_roles)),
+            ]
+            decoded["scene_memory_tokens"] = scene_memory.unsqueeze(1).expand(-1, self.config.scene_bank_size, -1)
+            decoded["belief_memory_tokens"] = belief_memory.unsqueeze(1).expand(-1, self.config.belief_bank_size, -1)
+            decoded["memory_tokens"] = torch.cat([decoded["scene_memory_tokens"], decoded["belief_memory_tokens"]], dim=1)
+            decoded["memory_token"] = decoded["memory_tokens"].mean(dim=1, keepdim=True)
+            decoded["uncertainty"] = decoded["uncertainty_field"].mean(dim=(-1, -2)).squeeze(1)
+            decoded["reocclusion_logit"] = decoded["reocclusion_field"].mean(dim=(-1, -2)).expand(-1, self.config.num_support_modes)
+            for key, value in decoded.items():
+                outputs.setdefault(key, []).append(value)
+
+        return {key: torch.stack(values, dim=1) for key, values in outputs.items()}

code/reveal_vla_bimanual/sim_reveal/dataset.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import pickle
 from pathlib import Path
 from typing import Any, Sequence
 
@@ -9,9 +10,10 @@ from torch.utils.data import Dataset
 
 import numpy as np
 
-from sim_reveal.procedural_envs import available_proxy_names, make_proxy_env, render_views_from_state
+from sim_reveal.procedural_envs import available_proxy_names, default_camera_matrices, make_proxy_env, render_views_from_state
 
 NOLEAK_PROXY_DATASET_VERSION = "reveal_proxy_v5_noleak_actionhist"
+RGBD_PROXY_DATASET_VERSION = "reveal_proxy_v6_rgbd_elastic_state"
 LEGACY_PRIVILEGED_RENDER_KEYS = frozenset(
     {
         "target_template",
@@ -44,6 +46,7 @@ def collect_teacher_dataset(
     rollout_horizon: int = 5,
     history_steps: int = 2,
     planner_candidates: int = 4,
+    dataset_version: str = NOLEAK_PROXY_DATASET_VERSION,
 ) -> dict[str, Any]:
     proxy_names = tuple(proxy_names or available_proxy_names())
     samples: list[dict[str, Any]] = []
@@ -91,7 +94,7 @@ def collect_teacher_dataset(
                     padded_history_actions.append(item["action"])
                 samples.append(
                     {
-                        "dataset_version": NOLEAK_PROXY_DATASET_VERSION,
+                        "dataset_version": dataset_version,
                         "proxy_name": proxy_name,
                         "episode_id": episode_idx,
                         "render_state": env.render_state(privileged_state),
@@ -103,10 +106,25 @@ def collect_teacher_dataset(
                         "persistence_horizon": privileged_state["persistence_horizon"].astype("float32"),
                         "disturbance_cost": float(privileged_state["disturbance_cost"]),
                         "belief_map": privileged_state["belief_map"].astype("float32"),
+                        "visibility_map": privileged_state["visibility_map"].astype("float32"),
+                        "clearance_map": privileged_state["clearance_map"].astype("float32"),
+                        "occluder_contact_map": privileged_state["occluder_contact_map"].astype("float32"),
+                        "grasp_affordance_map": privileged_state["grasp_affordance_map"].astype("float32"),
+                        "support_stability": float(privileged_state["support_stability"]),
+                        "support_stability_map": privileged_state["support_stability_map"].astype("float32"),
+                        "reocclusion_target": float(privileged_state["reocclusion_target"]),
+                        "reocclusion_map": privileged_state["reocclusion_map"].astype("float32"),
                         "rollout_support_mode": rollout["rollout_support_mode"].astype("int64"),
                         "rollout_corridor_feasible": rollout["rollout_corridor_feasible"].astype("float32"),
                         "rollout_persistence_horizon": rollout["rollout_persistence_horizon"].astype("float32"),
                         "rollout_disturbance_cost": rollout["rollout_disturbance_cost"].astype("float32"),
+                        "rollout_belief_map": rollout["rollout_belief_map"].astype("float32"),
+                        "rollout_visibility_map": rollout["rollout_visibility_map"].astype("float32"),
+                        "rollout_clearance_map": rollout["rollout_clearance_map"].astype("float32"),
+                        "rollout_support_stability": rollout["rollout_support_stability"].astype("float32"),
+                        "rollout_reocclusion_target": rollout["rollout_reocclusion_target"].astype("float32"),
+                        "rollout_occluder_contact_map": rollout["rollout_occluder_contact_map"].astype("float32"),
+                        "rollout_grasp_affordance_map": rollout["rollout_grasp_affordance_map"].astype("float32"),
                         "history_render_states": padded_history_render_states,
                         "history_proprio": np.stack(padded_history_proprio, axis=0).astype("float32")
                         if padded_history_proprio
@@ -138,7 +156,7 @@ def collect_teacher_dataset(
             "teacher_success": proxy_success / float(max(1, episodes_per_proxy)),
         }
     return {
-        "dataset_version": NOLEAK_PROXY_DATASET_VERSION,
+        "dataset_version": dataset_version,
         "resolution": resolution,
         "chunk_horizon": chunk_horizon,
         "rollout_horizon": rollout_horizon,
@@ -164,25 +182,46 @@ class RevealOfflineDataset(Dataset[dict[str, Any]]):
     def __init__(self, samples: Sequence[dict[str, Any]], resolution: int = 96) -> None:
         self.samples = list(samples)
         self.resolution = resolution
+        self._render_cache: dict[bytes, dict[str, np.ndarray]] = {}
+        self._item_cache: dict[int, dict[str, Any]] = {}
 
     def __len__(self) -> int:
         return len(self.samples)
 
-    def __getitem__(self, index: int) -> dict[str, Any]:
-        sample = self.samples[index]
-        _assert_noleak_sample(sample)
-        images = render_views_from_state(
+    def _render_cache_key(self, sample: dict[str, Any], render_state: dict[str, Any]) -> bytes:
+        include_depth = sample.get("dataset_version") == RGBD_PROXY_DATASET_VERSION
+        return pickle.dumps(
+            (sample["proxy_name"], self.resolution, include_depth, render_state),
+            protocol=4,
+        )
+
+    def _render_sample(self, sample: dict[str, Any], render_state: dict[str, Any]) -> dict[str, np.ndarray]:
+        cache_key = self._render_cache_key(sample, render_state)
+        cached = self._render_cache.get(cache_key)
+        if cached is not None:
+            return cached
+        include_depth = sample.get("dataset_version") == RGBD_PROXY_DATASET_VERSION
+        rendered = render_views_from_state(
             proxy_name=sample["proxy_name"],
-            render_state=sample["render_state"],
+            render_state=render_state,
             resolution=self.resolution,
+            include_depth=include_depth,
         )
+        self._render_cache[cache_key] = rendered
+        return rendered
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        cached_item = self._item_cache.get(index)
+        if cached_item is not None:
+            return cached_item
+        sample = self.samples[index]
+        _assert_noleak_sample(sample)
+        images = self._render_sample(sample, sample["render_state"])
         history_images = []
+        history_depths = []
+        history_depth_valid = []
         for history_state in sample.get("history_render_states", []):
-            rendered = render_views_from_state(
-                proxy_name=sample["proxy_name"],
-                render_state=history_state,
-                resolution=self.resolution,
-            )
+            rendered = self._render_sample(sample, history_state)
             history_images.append(
                 torch.stack(
                     [
@@ -193,6 +232,27 @@ class RevealOfflineDataset(Dataset[dict[str, Any]]):
                     dim=0,
                 )
             )
+            if sample.get("dataset_version") == RGBD_PROXY_DATASET_VERSION:
+                history_depths.append(
+                    torch.stack(
+                        [
+                            torch.from_numpy(rendered["front_depth"]),
+                            torch.from_numpy(rendered["wrist_left_depth"]),
+                            torch.from_numpy(rendered["wrist_right_depth"]),
+                        ],
+                        dim=0,
+                    )
+                )
+                history_depth_valid.append(
+                    torch.stack(
+                        [
+                            torch.from_numpy(rendered["front_depth_valid"]),
+                            torch.from_numpy(rendered["wrist_left_depth_valid"]),
+                            torch.from_numpy(rendered["wrist_right_depth_valid"]),
+                        ],
+                        dim=0,
+                    )
+                )
         stacked = torch.from_numpy(
             torch.stack(
                 [
@@ -207,9 +267,42 @@ class RevealOfflineDataset(Dataset[dict[str, Any]]):
             history_stacked = torch.stack(history_images, dim=0).permute(0, 1, 4, 2, 3).float() / 255.0
         else:
             history_stacked = torch.zeros((0, 3, 3, self.resolution, self.resolution), dtype=torch.float32)
-        return {
+        if sample.get("dataset_version") == RGBD_PROXY_DATASET_VERSION:
+            depths = torch.stack(
+                [
+                    torch.from_numpy(images["front_depth"]),
+                    torch.from_numpy(images["wrist_left_depth"]),
+                    torch.from_numpy(images["wrist_right_depth"]),
+                ],
+                dim=0,
+            ).unsqueeze(1).float()
+            depth_valid = torch.stack(
+                [
+                    torch.from_numpy(images["front_depth_valid"]),
+                    torch.from_numpy(images["wrist_left_depth_valid"]),
+                    torch.from_numpy(images["wrist_right_depth_valid"]),
+                ],
+                dim=0,
+            ).unsqueeze(1).float()
+            if history_depths:
+                history_depths_tensor = torch.stack(history_depths, dim=0).unsqueeze(2).float()
+                history_depth_valid_tensor = torch.stack(history_depth_valid, dim=0).unsqueeze(2).float()
+            else:
+                history_depths_tensor = torch.zeros((0, 3, 1, self.resolution, self.resolution), dtype=torch.float32)
+                history_depth_valid_tensor = torch.zeros((0, 3, 1, self.resolution, self.resolution), dtype=torch.float32)
+        else:
+            depths = torch.zeros((3, 1, self.resolution, self.resolution), dtype=torch.float32)
+            depth_valid = torch.zeros_like(depths)
+            history_depths_tensor = torch.zeros((0, 3, 1, self.resolution, self.resolution), dtype=torch.float32)
+            history_depth_valid_tensor = torch.zeros_like(history_depths_tensor)
+        camera_intrinsics, camera_extrinsics = default_camera_matrices()
+        item = {
             "images": stacked,
+            "depths": depths,
+            "depth_valid": depth_valid,
             "history_images": history_stacked,
+            "history_depths": history_depths_tensor,
+            "history_depth_valid": history_depth_valid_tensor,
             "history_proprio": torch.as_tensor(sample.get("history_proprio", []), dtype=torch.float32),
             "history_actions": torch.as_tensor(
                 sample.get(
@@ -218,6 +311,8 @@ class RevealOfflineDataset(Dataset[dict[str, Any]]):
                 ),
                 dtype=torch.float32,
             ),
+            "camera_intrinsics": torch.as_tensor(camera_intrinsics, dtype=torch.float32),
+            "camera_extrinsics": torch.as_tensor(camera_extrinsics, dtype=torch.float32),
             "proprio": torch.as_tensor(sample["proprio"], dtype=torch.float32),
             "texts": sample["language_goal"],
             "action_chunk": torch.as_tensor(sample["action_chunk"], dtype=torch.float32),
@@ -226,15 +321,37 @@ class RevealOfflineDataset(Dataset[dict[str, Any]]):
             "persistence_horizon": torch.as_tensor(sample["persistence_horizon"], dtype=torch.float32),
             "disturbance_cost": torch.as_tensor(sample["disturbance_cost"], dtype=torch.float32),
             "belief_map": torch.as_tensor(sample["belief_map"], dtype=torch.float32).unsqueeze(0),
+            "visibility_map": torch.as_tensor(sample.get("visibility_map", np.zeros((32, 32), dtype=np.float32)), dtype=torch.float32).unsqueeze(0),
+            "clearance_map": torch.as_tensor(sample.get("clearance_map", np.zeros((2, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "occluder_contact_map": torch.as_tensor(sample.get("occluder_contact_map", np.zeros((32, 32), dtype=np.float32)), dtype=torch.float32).unsqueeze(0),
+            "grasp_affordance_map": torch.as_tensor(sample.get("grasp_affordance_map", np.zeros((32, 32), dtype=np.float32)), dtype=torch.float32).unsqueeze(0),
+            "support_stability": torch.as_tensor(sample.get("support_stability", 0.0), dtype=torch.float32),
+            "support_stability_map": torch.as_tensor(sample.get("support_stability_map", np.zeros((32, 32), dtype=np.float32)), dtype=torch.float32).unsqueeze(0),
+            "reocclusion_target": torch.as_tensor(sample.get("reocclusion_target", 0.0), dtype=torch.float32),
+            "reocclusion_map": torch.as_tensor(sample.get("reocclusion_map", np.zeros((32, 32), dtype=np.float32)), dtype=torch.float32).unsqueeze(0),
             "rollout_support_mode": torch.as_tensor(sample["rollout_support_mode"], dtype=torch.long),
             "rollout_corridor_feasible": torch.as_tensor(sample["rollout_corridor_feasible"], dtype=torch.float32),
             "rollout_persistence_horizon": torch.as_tensor(sample["rollout_persistence_horizon"], dtype=torch.float32),
             "rollout_disturbance_cost": torch.as_tensor(sample["rollout_disturbance_cost"], dtype=torch.float32),
+            "rollout_belief_map": torch.as_tensor(sample.get("rollout_belief_map", np.zeros((0, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "rollout_visibility_map": torch.as_tensor(sample.get("rollout_visibility_map", np.zeros((0, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "rollout_clearance_map": torch.as_tensor(sample.get("rollout_clearance_map", np.zeros((0, 2, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "rollout_support_stability": torch.as_tensor(sample.get("rollout_support_stability", np.zeros((0,), dtype=np.float32)), dtype=torch.float32),
+            "rollout_reocclusion_target": torch.as_tensor(sample.get("rollout_reocclusion_target", np.zeros((0,), dtype=np.float32)), dtype=torch.float32),
+            "rollout_occluder_contact_map": torch.as_tensor(sample.get("rollout_occluder_contact_map", np.zeros((0, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "rollout_grasp_affordance_map": torch.as_tensor(sample.get("rollout_grasp_affordance_map", np.zeros((0, 32, 32), dtype=np.float32)), dtype=torch.float32),
             "candidate_action_chunks": torch.as_tensor(sample["candidate_action_chunks"], dtype=torch.float32),
             "candidate_rollout_support_mode": torch.as_tensor(sample["candidate_rollout_support_mode"], dtype=torch.long),
             "candidate_rollout_corridor_feasible": torch.as_tensor(sample["candidate_rollout_corridor_feasible"], dtype=torch.float32),
             "candidate_rollout_persistence_horizon": torch.as_tensor(sample["candidate_rollout_persistence_horizon"], dtype=torch.float32),
             "candidate_rollout_disturbance_cost": torch.as_tensor(sample["candidate_rollout_disturbance_cost"], dtype=torch.float32),
+            "candidate_rollout_belief_map": torch.as_tensor(sample.get("candidate_rollout_belief_map", np.zeros((0, 0, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "candidate_rollout_visibility_map": torch.as_tensor(sample.get("candidate_rollout_visibility_map", np.zeros((0, 0, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "candidate_rollout_clearance_map": torch.as_tensor(sample.get("candidate_rollout_clearance_map", np.zeros((0, 0, 2, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "candidate_rollout_support_stability": torch.as_tensor(sample.get("candidate_rollout_support_stability", np.zeros((0, 0), dtype=np.float32)), dtype=torch.float32),
+            "candidate_rollout_reocclusion_target": torch.as_tensor(sample.get("candidate_rollout_reocclusion_target", np.zeros((0, 0), dtype=np.float32)), dtype=torch.float32),
+            "candidate_rollout_occluder_contact_map": torch.as_tensor(sample.get("candidate_rollout_occluder_contact_map", np.zeros((0, 0, 32, 32), dtype=np.float32)), dtype=torch.float32),
+            "candidate_rollout_grasp_affordance_map": torch.as_tensor(sample.get("candidate_rollout_grasp_affordance_map", np.zeros((0, 0, 32, 32), dtype=np.float32)), dtype=torch.float32),
             "candidate_retrieval_success": torch.as_tensor(sample["candidate_retrieval_success"], dtype=torch.float32),
             "candidate_final_disturbance_cost": torch.as_tensor(sample["candidate_final_disturbance_cost"], dtype=torch.float32),
             "candidate_reocclusion_rate": torch.as_tensor(sample["candidate_reocclusion_rate"], dtype=torch.float32),
@@ -244,6 +361,8 @@ class RevealOfflineDataset(Dataset[dict[str, Any]]):
             "proxy_name": sample["proxy_name"],
             "episode_id": sample["episode_id"],
         }
+        self._item_cache[index] = item
+        return item
 
 
 def dataset_from_bundle(dataset_bundle: dict[str, Any], resolution: int | None = None) -> RevealOfflineDataset:

code/reveal_vla_bimanual/sim_reveal/procedural_envs.py

@@ -83,6 +83,26 @@ PROXY_GOALS = {
 }
 
 
+def default_camera_matrices() -> tuple[np.ndarray, np.ndarray]:
+    intrinsics = np.asarray(
+        [
+            [[140.0, 0.0, 48.0], [0.0, 140.0, 48.0], [0.0, 0.0, 1.0]],
+            [[135.0, 0.0, 48.0], [0.0, 135.0, 48.0], [0.0, 0.0, 1.0]],
+            [[135.0, 0.0, 48.0], [0.0, 135.0, 48.0], [0.0, 0.0, 1.0]],
+        ],
+        dtype=np.float32,
+    )
+    extrinsics = np.asarray(
+        [
+            np.eye(4, dtype=np.float32),
+            [[1.0, 0.0, 0.0, -0.18], [0.0, 1.0, 0.0, 0.04], [0.0, 0.0, 1.0, 0.10], [0.0, 0.0, 0.0, 1.0]],
+            [[1.0, 0.0, 0.0, 0.18], [0.0, 1.0, 0.0, 0.04], [0.0, 0.0, 1.0, 0.10], [0.0, 0.0, 0.0, 1.0]],
+        ],
+        dtype=np.float32,
+    )
+    return intrinsics, extrinsics
+
+
 def available_proxy_names() -> tuple[str, ...]:
     return tuple(PROXY_CONFIGS.keys())
 
@@ -285,6 +305,57 @@ class ProceduralRevealEnv:
         belief *= visibility
         return belief.astype(np.float32)
 
+    def _visibility_map(self, visibility: float) -> np.ndarray:
+        belief = self._belief_map(visibility)
+        gradient = np.linspace(0.65, 1.0, belief.shape[0], dtype=np.float32).reshape(-1, 1)
+        return np.clip(belief * gradient, 0.0, 1.0).astype(np.float32)
+
+    def _clearance_map(self, visibility: float) -> np.ndarray:
+        side = 32
+        x = np.linspace(0.0, 1.0, side, dtype=np.float32)
+        y = np.linspace(0.0, 1.0, side, dtype=np.float32)
+        yy, xx = np.meshgrid(y, x, indexing="ij")
+        corridor_width = np.clip(0.05 + 0.18 * self.opening - 0.10 * self.disturbance, 0.01, 0.28)
+        corridor = np.exp(-(((xx - self.target_center) ** 2) / max(1e-5, corridor_width**2)))
+        vertical = np.exp(-(((yy - (0.72 - 0.25 * self.target_depth)) ** 2) / 0.03))
+        left = np.clip(corridor * vertical * visibility * (0.92 - 0.15 * self.disturbance), 0.0, 1.0)
+        right = np.clip(corridor * vertical * visibility * (0.88 - 0.10 * self.disturbance), 0.0, 1.0)
+        return np.stack([left, right], axis=0).astype(np.float32)
+
+    def _occluder_contact_map(self) -> np.ndarray:
+        side = 32
+        x = np.linspace(0.0, 1.0, side, dtype=np.float32)
+        y = np.linspace(0.0, 1.0, side, dtype=np.float32)
+        yy, xx = np.meshgrid(y, x, indexing="ij")
+        gap_width = np.clip(0.03 + 0.16 * self.opening, 0.03, 0.24)
+        left_band = np.exp(-(((xx - (self.target_center - gap_width)) ** 2) / 0.0025))
+        right_band = np.exp(-(((xx - (self.target_center + gap_width)) ** 2) / 0.0025))
+        support = np.exp(-(((yy - 0.55) ** 2) / 0.04))
+        return np.clip((left_band + right_band) * support, 0.0, 1.0).astype(np.float32)
+
+    def _support_stability(self) -> float:
+        base = 1.0 - 0.45 * self.disturbance - 0.10 * max(0.0, self.opening - self.dynamics.desired_opening)
+        if self._current_support_mode() == self.dynamics.preferred_mode:
+            base += 0.08
+        return float(np.clip(base, 0.0, 1.0))
+
+    def _support_stability_map(self) -> np.ndarray:
+        return np.full((32, 32), self._support_stability(), dtype=np.float32)
+
+    def _reocclusion_target(self, persistence: np.ndarray) -> float:
+        current_mode = int(self._current_support_mode())
+        horizon_ratio = persistence[current_mode] / float(max(1, self.rollout_horizon))
+        return float(np.clip(1.0 - horizon_ratio + 0.35 * self.disturbance, 0.0, 1.0))
+
+    def _grasp_affordance_map(
+        self,
+        belief_map: np.ndarray,
+        visibility_map: np.ndarray,
+        clearance_map: np.ndarray,
+    ) -> np.ndarray:
+        combined = belief_map * visibility_map * clearance_map.mean(axis=0)
+        return np.clip(combined * (1.0 - 0.35 * self.disturbance), 0.0, 1.0).astype(np.float32)
+
     def get_privileged_state(self) -> dict[str, Any]:
         support_mode = int(self._current_support_mode())
         corridor = np.stack(
@@ -294,12 +365,29 @@ class ProceduralRevealEnv:
         persistence = np.asarray([self._persistence_for_mode(mode) for mode in SupportMode], dtype=np.float32)
         visibility = self._visibility()
         disturbance_cost = float(np.clip(self.disturbance + 0.08 * max(0.0, self.opening - self.dynamics.desired_opening), 0.0, 1.0))
+        belief_map = self._belief_map(visibility)
+        visibility_map = self._visibility_map(visibility)
+        clearance_map = self._clearance_map(visibility)
+        occluder_contact_map = self._occluder_contact_map()
+        support_stability = self._support_stability()
+        support_stability_map = self._support_stability_map()
+        reocclusion_target = self._reocclusion_target(persistence)
+        reocclusion_map = np.full((32, 32), reocclusion_target, dtype=np.float32)
+        grasp_affordance_map = self._grasp_affordance_map(belief_map, visibility_map, clearance_map)
         return {
             "support_mode": support_mode,
             "corridor_feasible": corridor,
             "persistence_horizon": persistence,
             "disturbance_cost": disturbance_cost,
-            "belief_map": self._belief_map(visibility),
+            "belief_map": belief_map,
+            "visibility_map": visibility_map,
+            "clearance_map": clearance_map,
+            "occluder_contact_map": occluder_contact_map,
+            "grasp_affordance_map": grasp_affordance_map,
+            "support_stability": support_stability,
+            "support_stability_map": support_stability_map,
+            "reocclusion_target": reocclusion_target,
+            "reocclusion_map": reocclusion_map,
             "visibility": visibility,
             "retrieval_success": bool(self.retrieved),
             "target_template": self.target_template,
@@ -335,12 +423,18 @@ class ProceduralRevealEnv:
             render_state=render_state,
             resolution=self.resolution,
             num_templates=self.num_templates,
+            include_depth=True,
         )
+        camera_intrinsics, camera_extrinsics = default_camera_matrices()
         return {
             "images": np.stack([images[camera] for camera in self.camera_names], axis=0),
+            "depths": np.stack([images[f"{camera}_depth"] for camera in self.camera_names], axis=0)[:, None, :, :],
+            "depth_valid": np.stack([images[f"{camera}_depth_valid"] for camera in self.camera_names], axis=0)[:, None, :, :],
             "proprio": self._proprio(privileged_state),
             "text": PROXY_GOALS[self.proxy_name],
             "camera_names": self.camera_names,
+            "camera_intrinsics": camera_intrinsics,
+            "camera_extrinsics": camera_extrinsics,
         }
 
     def teacher_action(self) -> np.ndarray:
@@ -385,6 +479,13 @@ class ProceduralRevealEnv:
         rollout_corridor = []
         rollout_persistence = []
         rollout_disturbance = []
+        rollout_belief = []
+        rollout_visibility = []
+        rollout_clearance = []
+        rollout_support_stability = []
+        rollout_reocclusion = []
+        rollout_occluder_contact = []
+        rollout_grasp_affordance = []
         for step in range(chunk_horizon):
             action = self.teacher_action()
             action_chunk.append(action)
@@ -394,21 +495,43 @@ class ProceduralRevealEnv:
                 rollout_corridor.append(privileged_state["corridor_feasible"])
                 rollout_persistence.append(privileged_state["persistence_horizon"])
                 rollout_disturbance.append(privileged_state["disturbance_cost"])
+                rollout_belief.append(privileged_state["belief_map"])
+                rollout_visibility.append(privileged_state["visibility_map"])
+                rollout_clearance.append(privileged_state["clearance_map"])
+                rollout_support_stability.append(privileged_state["support_stability"])
+                rollout_reocclusion.append(privileged_state["reocclusion_target"])
+                rollout_occluder_contact.append(privileged_state["occluder_contact_map"])
+                rollout_grasp_affordance.append(privileged_state["grasp_affordance_map"])
             if terminated or truncated:
                 break
         while len(action_chunk) < chunk_horizon:
             action_chunk.append(np.zeros((14,), dtype=np.float32))
         while len(rollout_support_mode) < rollout_horizon:
+            current = self.get_privileged_state()
             rollout_support_mode.append(int(self._current_support_mode()))
-            rollout_corridor.append(self.get_privileged_state()["corridor_feasible"])
-            rollout_persistence.append(self.get_privileged_state()["persistence_horizon"])
-            rollout_disturbance.append(self.get_privileged_state()["disturbance_cost"])
+            rollout_corridor.append(current["corridor_feasible"])
+            rollout_persistence.append(current["persistence_horizon"])
+            rollout_disturbance.append(current["disturbance_cost"])
+            rollout_belief.append(current["belief_map"])
+            rollout_visibility.append(current["visibility_map"])
+            rollout_clearance.append(current["clearance_map"])
+            rollout_support_stability.append(current["support_stability"])
+            rollout_reocclusion.append(current["reocclusion_target"])
+            rollout_occluder_contact.append(current["occluder_contact_map"])
+            rollout_grasp_affordance.append(current["grasp_affordance_map"])
         self.restore_state(snapshot)
         return np.stack(action_chunk, axis=0).astype(np.float32), {
             "rollout_support_mode": np.asarray(rollout_support_mode, dtype=np.int64),
             "rollout_corridor_feasible": np.asarray(rollout_corridor, dtype=np.float32),
             "rollout_persistence_horizon": np.asarray(rollout_persistence, dtype=np.float32),
             "rollout_disturbance_cost": np.asarray(rollout_disturbance, dtype=np.float32),
+            "rollout_belief_map": np.asarray(rollout_belief, dtype=np.float32),
+            "rollout_visibility_map": np.asarray(rollout_visibility, dtype=np.float32),
+            "rollout_clearance_map": np.asarray(rollout_clearance, dtype=np.float32),
+            "rollout_support_stability": np.asarray(rollout_support_stability, dtype=np.float32),
+            "rollout_reocclusion_target": np.asarray(rollout_reocclusion, dtype=np.float32),
+            "rollout_occluder_contact_map": np.asarray(rollout_occluder_contact, dtype=np.float32),
+            "rollout_grasp_affordance_map": np.asarray(rollout_grasp_affordance, dtype=np.float32),
         }
 
     def evaluate_action_chunk(
@@ -422,6 +545,13 @@ class ProceduralRevealEnv:
         rollout_corridor: list[np.ndarray] = []
         rollout_persistence: list[np.ndarray] = []
         rollout_disturbance: list[float] = []
+        rollout_belief: list[np.ndarray] = []
+        rollout_visibility: list[np.ndarray] = []
+        rollout_clearance: list[np.ndarray] = []
+        rollout_support_stability: list[float] = []
+        rollout_reocclusion: list[float] = []
+        rollout_occluder_contact: list[np.ndarray] = []
+        rollout_grasp_affordance: list[np.ndarray] = []
         corridor_open_trace = [float(self.get_privileged_state()["corridor_feasible"][self._current_support_mode()].any())]
         visibility_trace = [float(self.get_privileged_state()["visibility"])]
         terminated = False
@@ -434,6 +564,13 @@ class ProceduralRevealEnv:
                 rollout_corridor.append(privileged_state["corridor_feasible"].astype(np.float32))
                 rollout_persistence.append(privileged_state["persistence_horizon"].astype(np.float32))
                 rollout_disturbance.append(float(privileged_state["disturbance_cost"]))
+                rollout_belief.append(privileged_state["belief_map"].astype(np.float32))
+                rollout_visibility.append(privileged_state["visibility_map"].astype(np.float32))
+                rollout_clearance.append(privileged_state["clearance_map"].astype(np.float32))
+                rollout_support_stability.append(float(privileged_state["support_stability"]))
+                rollout_reocclusion.append(float(privileged_state["reocclusion_target"]))
+                rollout_occluder_contact.append(privileged_state["occluder_contact_map"].astype(np.float32))
+                rollout_grasp_affordance.append(privileged_state["grasp_affordance_map"].astype(np.float32))
             corridor_open_trace.append(float(privileged_state["corridor_feasible"][privileged_state["support_mode"]].any()))
             visibility_trace.append(float(privileged_state["visibility"]))
             if terminated or truncated:
@@ -444,6 +581,13 @@ class ProceduralRevealEnv:
             rollout_corridor.append(current["corridor_feasible"].astype(np.float32))
             rollout_persistence.append(current["persistence_horizon"].astype(np.float32))
             rollout_disturbance.append(float(current["disturbance_cost"]))
+            rollout_belief.append(current["belief_map"].astype(np.float32))
+            rollout_visibility.append(current["visibility_map"].astype(np.float32))
+            rollout_clearance.append(current["clearance_map"].astype(np.float32))
+            rollout_support_stability.append(float(current["support_stability"]))
+            rollout_reocclusion.append(float(current["reocclusion_target"]))
+            rollout_occluder_contact.append(current["occluder_contact_map"].astype(np.float32))
+            rollout_grasp_affordance.append(current["grasp_affordance_map"].astype(np.float32))
         final_state = self.get_privileged_state()
         reocclusion = float(
             np.logical_and(
@@ -456,6 +600,13 @@ class ProceduralRevealEnv:
             "rollout_corridor_feasible": np.asarray(rollout_corridor, dtype=np.float32),
             "rollout_persistence_horizon": np.asarray(rollout_persistence, dtype=np.float32),
             "rollout_disturbance_cost": np.asarray(rollout_disturbance, dtype=np.float32),
+            "rollout_belief_map": np.asarray(rollout_belief, dtype=np.float32),
+            "rollout_visibility_map": np.asarray(rollout_visibility, dtype=np.float32),
+            "rollout_clearance_map": np.asarray(rollout_clearance, dtype=np.float32),
+            "rollout_support_stability": np.asarray(rollout_support_stability, dtype=np.float32),
+            "rollout_reocclusion_target": np.asarray(rollout_reocclusion, dtype=np.float32),
+            "rollout_occluder_contact_map": np.asarray(rollout_occluder_contact, dtype=np.float32),
+            "rollout_grasp_affordance_map": np.asarray(rollout_grasp_affordance, dtype=np.float32),
             "retrieval_success": float(final_state["retrieval_success"]),
             "final_disturbance_cost": float(final_state["disturbance_cost"]),
             "reocclusion_rate": reocclusion,
@@ -493,6 +644,27 @@ class ProceduralRevealEnv:
             "candidate_rollout_disturbance_cost": np.stack(
                 [item["rollout_disturbance_cost"] for item in outcomes], axis=0
             ).astype(np.float32),
+            "candidate_rollout_belief_map": np.stack(
+                [item["rollout_belief_map"] for item in outcomes], axis=0
+            ).astype(np.float32),
+            "candidate_rollout_visibility_map": np.stack(
+                [item["rollout_visibility_map"] for item in outcomes], axis=0
+            ).astype(np.float32),
+            "candidate_rollout_clearance_map": np.stack(
+                [item["rollout_clearance_map"] for item in outcomes], axis=0
+            ).astype(np.float32),
+            "candidate_rollout_support_stability": np.stack(
+                [item["rollout_support_stability"] for item in outcomes], axis=0
+            ).astype(np.float32),
+            "candidate_rollout_reocclusion_target": np.stack(
+                [item["rollout_reocclusion_target"] for item in outcomes], axis=0
+            ).astype(np.float32),
+            "candidate_rollout_occluder_contact_map": np.stack(
+                [item["rollout_occluder_contact_map"] for item in outcomes], axis=0
+            ).astype(np.float32),
+            "candidate_rollout_grasp_affordance_map": np.stack(
+                [item["rollout_grasp_affordance_map"] for item in outcomes], axis=0
+            ).astype(np.float32),
             "candidate_retrieval_success": np.asarray([item["retrieval_success"] for item in outcomes], dtype=np.float32),
             "candidate_final_disturbance_cost": np.asarray(
                 [item["final_disturbance_cost"] for item in outcomes], dtype=np.float32
@@ -587,6 +759,7 @@ def render_views_from_state(
     render_state: dict[str, Any],
     resolution: int,
     num_templates: int = 32,
+    include_depth: bool = False,
 ) -> dict[str, np.ndarray]:
     dynamics = PROXY_DYNAMICS[proxy_name]
     opening = float(render_state["opening"])
@@ -668,8 +841,40 @@ def render_views_from_state(
     wrist_right[..., 2] = np.clip(wrist_right[..., 2] + 0.08 * step_fraction + 0.06 * right_band, 0.0, 1.0)
     wrist_right = np.clip(wrist_right, 0.0, 1.0)
 
-    return {
+    outputs = {
         "front": (front * 255.0).astype(np.uint8),
         "wrist_left": (wrist_left * 255.0).astype(np.uint8),
         "wrist_right": (wrist_right * 255.0).astype(np.uint8),
     }
+    if not include_depth:
+        return outputs
+
+    front_depth = np.clip(0.25 + 0.40 * target_depth + 0.15 * disturbance + 0.10 * (1.0 - visibility), 0.0, 1.0)
+    target_depth_map = np.clip(0.10 + 0.55 * target_depth, 0.0, 1.0)
+    occluder_depth = np.clip(0.30 + 0.20 * disturbance + 0.10 * (1.0 - opening), 0.0, 1.0)
+    front_depth_map = np.full((height, width), front_depth, dtype=np.float32)
+    front_depth_map[gap_mask] = np.minimum(front_depth_map[gap_mask], occluder_depth)
+    front_depth_map[target_mask] = np.minimum(front_depth_map[target_mask], target_depth_map)
+
+    wrist_left_depth = np.clip(0.35 + 0.25 * target_depth + 0.10 * disturbance, 0.0, 1.0)
+    wrist_left_depth_map = np.full((height, width), wrist_left_depth, dtype=np.float32)
+    wrist_left_depth_map[left_open] = np.minimum(wrist_left_depth_map[left_open], 0.22 + 0.25 * target_depth)
+    wrist_left_depth_map[target_mask] = np.minimum(wrist_left_depth_map[target_mask], target_depth_map)
+
+    wrist_right_depth = np.clip(0.35 + 0.20 * target_depth + 0.12 * disturbance, 0.0, 1.0)
+    wrist_right_depth_map = np.full((height, width), wrist_right_depth, dtype=np.float32)
+    right_focus = (right_band * right_clear) > 0.15
+    wrist_right_depth_map[right_focus] = np.minimum(wrist_right_depth_map[right_focus], 0.20 + 0.25 * target_depth)
+    wrist_right_depth_map[target_mask] = np.minimum(wrist_right_depth_map[target_mask], target_depth_map)
+
+    outputs.update(
+        {
+            "front_depth": front_depth_map.astype(np.float32),
+            "wrist_left_depth": wrist_left_depth_map.astype(np.float32),
+            "wrist_right_depth": wrist_right_depth_map.astype(np.float32),
+            "front_depth_valid": np.ones((height, width), dtype=np.float32),
+            "wrist_left_depth_valid": np.ones((height, width), dtype=np.float32),
+            "wrist_right_depth_valid": np.ones((height, width), dtype=np.float32),
+        }
+    )
+    return outputs