Add files using upload-large-folder tool

code/reveal_vla_bimanual/eval/ablations.py

@@ -1,8 +1,7 @@
 MANDATORY_ABLATIONS: tuple[str, ...] = (
-    "no_reveal_state_head",
+    "no_interaction_head",
     "no_world_model",
-    "no_planner_reranking",
-    "no_support_mode_conditioning",
-    "no_wrist_cameras",
-    "no_global_camera",
+    "no_planner",
+    "no_role_tokens",
+    "short_history",
 )

code/reveal_vla_bimanual/eval/metrics.py

@@ -16,6 +16,14 @@ class BenchmarkMetrics:
     disturbance_cost: float | None = None
 
 
+@dataclass
+class PlannerDiagnostics:
+    top1_accuracy: float
+    regret: float
+    risk_calibration_mse: float
+    role_collapse_rate: float
+
+
 def mean_success(per_task_success: dict[str, float]) -> float:
     if not per_task_success:
         return 0.0
@@ -50,3 +58,50 @@ def mean_disturbance_cost(values: np.ndarray) -> float:
     if values.size == 0:
         return 0.0
     return float(values.mean())
+
+
+def planner_top1_accuracy(pred_scores: np.ndarray, oracle_utility: np.ndarray) -> float:
+    pred_scores = np.asarray(pred_scores)
+    oracle_utility = np.asarray(oracle_utility)
+    if pred_scores.size == 0:
+        return 0.0
+    return float((pred_scores.argmax(axis=-1) == oracle_utility.argmax(axis=-1)).mean())
+
+
+def planner_regret(selected_indices: np.ndarray, oracle_utility: np.ndarray) -> float:
+    selected_indices = np.asarray(selected_indices, dtype=np.int64)
+    oracle_utility = np.asarray(oracle_utility, dtype=np.float32)
+    if oracle_utility.size == 0:
+        return 0.0
+    batch_index = np.arange(selected_indices.shape[0])
+    selected = oracle_utility[batch_index, selected_indices]
+    oracle = oracle_utility.max(axis=-1)
+    return float((oracle - selected).mean())
+
+
+def risk_calibration_mse(predicted_risk: np.ndarray, realized_risk: np.ndarray) -> float:
+    predicted_risk = np.asarray(predicted_risk, dtype=np.float32)
+    realized_risk = np.asarray(realized_risk, dtype=np.float32)
+    if predicted_risk.size == 0:
+        return 0.0
+    return float(np.mean((predicted_risk - realized_risk) ** 2))
+
+
+def role_collapse_rate(
+    action_chunks: np.ndarray,
+    arm_role_logits: np.ndarray | None = None,
+    action_threshold: float = 1e-2,
+    role_threshold: float = 0.1,
+) -> float:
+    action_chunks = np.asarray(action_chunks, dtype=np.float32)
+    right_actions = action_chunks[..., :7]
+    left_actions = action_chunks[..., 7:]
+    action_gap = np.mean(np.abs(right_actions - left_actions), axis=(-1, -2))
+    collapsed = action_gap <= action_threshold
+    if arm_role_logits is not None:
+        arm_role_logits = np.asarray(arm_role_logits, dtype=np.float32)
+        role_probs = np.exp(arm_role_logits - arm_role_logits.max(axis=-1, keepdims=True))
+        role_probs = role_probs / np.clip(role_probs.sum(axis=-1, keepdims=True), 1e-6, None)
+        role_gap = np.mean(np.abs(role_probs[..., 0, :] - role_probs[..., 1, :]), axis=-1)
+        collapsed = np.logical_or(collapsed, role_gap <= role_threshold)
+    return float(collapsed.mean())

code/reveal_vla_bimanual/eval/run_proxy_diagnostics.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+from eval.metrics import planner_regret, planner_top1_accuracy, risk_calibration_mse, role_collapse_rate
+from eval.run_reveal_benchmark import load_model
+from sim_reveal.dataset import dataset_from_bundle, load_teacher_dataset
+
+
+def _move_batch_to_device(batch: dict[str, Any], device: torch.device) -> dict[str, Any]:
+    moved = {}
+    for key, value in batch.items():
+        if isinstance(value, Tensor):
+            moved[key] = value.to(device)
+        else:
+            moved[key] = value
+    return moved
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--dataset", required=True)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--num-workers", type=int, default=0)
+    parser.add_argument("--output-dir", required=True)
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model, _ = load_model(args.checkpoint, device=device)
+    bundle = load_teacher_dataset(args.dataset)
+    dataset = dataset_from_bundle(bundle, resolution=int(bundle["resolution"]))
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+
+    score_batches: list[np.ndarray] = []
+    utility_batches: list[np.ndarray] = []
+    best_index_batches: list[np.ndarray] = []
+    risk_batches: list[np.ndarray] = []
+    realized_risk_batches: list[np.ndarray] = []
+    collapse_batches: list[float] = []
+
+    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"),
+                plan=True,
+                candidate_chunks_override=moved["candidate_action_chunks"],
+            )
+            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()
+            )
+            selected_chunk = outputs["planned_chunk"].detach().cpu().numpy()[:, None]
+            role_logits = None
+            if outputs.get("interaction_state") is not None:
+                role_logits = outputs["interaction_state"]["arm_role_logits"].detach().cpu().numpy()[:, None]
+            collapse_batches.append(role_collapse_rate(selected_chunk, role_logits))
+
+    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)
+    selected_indices = (
+        np.concatenate(best_index_batches, axis=0) if best_index_batches else np.zeros((0,), dtype=np.int64)
+    )
+    predicted_risk = np.concatenate(risk_batches, axis=0) if risk_batches else np.zeros((0, 0), dtype=np.float32)
+    realized_risk = (
+        np.concatenate(realized_risk_batches, axis=0) if realized_risk_batches else np.zeros((0, 0), dtype=np.float32)
+    )
+
+    diagnostics = {
+        "planner_top1_accuracy": planner_top1_accuracy(scores, utility),
+        "planner_regret": planner_regret(selected_indices, 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,
+        "num_samples": int(scores.shape[0]),
+    }
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    (output_dir / "proxy_diagnostics.json").write_text(json.dumps(diagnostics, indent=2), encoding="utf-8")
+    print(json.dumps(diagnostics, indent=2))
+
+
+if __name__ == "__main__":
+    main()

code/reveal_vla_bimanual/eval/run_reveal_benchmark.py

@@ -49,7 +49,7 @@ def _trainer_config_from_dict(cfg: dict[str, Any]) -> TrainerConfig:
 
 
 def load_model(checkpoint_path: str | Path, device: torch.device) -> tuple[torch.nn.Module, dict[str, Any]]:
-    checkpoint = torch.load(Path(checkpoint_path), map_location="cpu")
+    checkpoint = torch.load(Path(checkpoint_path), map_location="cpu", weights_only=False)
     policy_config = _policy_config_from_dict(checkpoint["policy_config"])
     trainer_config = _trainer_config_from_dict(checkpoint["trainer_config"])
     model = build_policy(policy_config, trainer_config).to(device)
@@ -112,6 +112,22 @@ def select_chunk(
         "proprio": batch["proprio"],
         "texts": batch["texts"],
     }
+    if hasattr(model, "interaction_head"):
+        outputs = model(
+            **forward_kwargs,
+            plan=(ablation not in {"no_world_model", "no_interaction_head"}),
+            support_mode_conditioning=True,
+            use_interaction_head=(ablation != "no_interaction_head"),
+            use_role_tokens=(ablation != "no_role_tokens"),
+            history_steps_override=(2 if ablation == "short_history" else None),
+        )
+        if ablation == "no_planner":
+            if "candidate_chunks" in outputs:
+                return outputs["candidate_chunks"][:, 0], outputs
+            return outputs["action_mean"], outputs
+        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, "reveal_head"):
         if ablation == "no_world_model":
             outputs = model(**forward_kwargs, plan=False)
@@ -181,10 +197,13 @@ def evaluate_model(
                 episode_visibility.append(float(privileged_state["visibility"]))
                 episode_corridor.append(float(privileged_state["corridor_feasible"][privileged_state["support_mode"]].any()))
                 episode_disturbance.append(float(privileged_state["disturbance_cost"]))
-                if "reveal_state" in outputs and ablation != "no_reveal_state_head":
+                state_output = outputs.get("interaction_state")
+                if state_output is None:
+                    state_output = outputs.get("reveal_state")
+                if state_output is not None and ablation != "no_interaction_head":
                     persistence_errors.append(
                         persistence_horizon_mae(
-                            outputs["reveal_state"]["persistence_horizon"][0].detach().cpu().numpy(),
+                            state_output["persistence_horizon"][0].detach().cpu().numpy(),
                             privileged_state["persistence_horizon"],
                         )
                     )

code/reveal_vla_bimanual/eval/run_rlbench_rollout_eval.py

@@ -7,13 +7,12 @@ from typing import Any, Sequence
 
 import numpy as np
 import torch
-from helpers.observation_utils import create_obs_config
-from omegaconf import OmegaConf
-from rlbench.action_modes.action_mode import BimanualMoveArmThenGripper
-from rlbench.action_modes.arm_action_modes import BimanualEndEffectorPoseViaPlanning
-from rlbench.action_modes.gripper_action_modes import BimanualDiscrete
+from helpers.utils import create_obs_config
+from rlbench.action_modes.action_mode import MoveArmThenGripper2Robots
+from rlbench.action_modes.arm_action_modes import EndEffectorPoseViaPlanning2Robots
+from rlbench.action_modes.gripper_action_modes import Discrete2Robots
 from rlbench.backend.utils import task_file_to_task_class
-from rlbench.environment import Environment
+from rlbench.environments_two_robots import Environment2Robots
 
 from models.action_decoder import ChunkDecoderConfig
 from models.backbones import FrozenVLBackboneConfig
@@ -23,8 +22,13 @@ from models.planner import PlannerConfig
 from models.policy import PolicyConfig
 from models.reveal_head import RevealHeadConfig
 from models.world_model import RevealWMConfig
-from sim_rlbench.dataset import absolute_action_from_delta, bimanual_proprio_from_obs, stack_live_rgb_obs
-from train.trainer import TrainerConfig, build_policy
+from sim_rlbench.camera_spec import default_three_camera_spec
+from sim_rlbench.dataset import (
+    bimanual_proprio_from_obs,
+    single_arm_absolute_action_from_delta,
+    stack_live_rgb_obs,
+)
+from train.trainer import TrainerConfig, build_policy, planner_enabled, policy_supports_planning
 
 
 def _policy_config_from_checkpoint(checkpoint: dict[str, Any]) -> PolicyConfig:
@@ -48,6 +52,19 @@ def _episode_language_goal(descriptions: Sequence[str]) -> str:
     return str(descriptions[0]) if descriptions else ""
 
 
+def _step_bimanual_chunk(task: Any, obs: Any, delta_action: np.ndarray) -> tuple[Any, float, bool]:
+    total_reward = 0.0
+    done = False
+    next_obs = obs
+    for arm_name in ("right", "left"):
+        env_action = single_arm_absolute_action_from_delta(next_obs, delta_action, arm_name, ignore_collisions=True)
+        next_obs, reward, done = task.step(env_action, arm_name)
+        total_reward += float(reward)
+        if reward >= 1.0 or done:
+            break
+    return next_obs, total_reward, done
+
+
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--checkpoint", required=True)
@@ -71,55 +88,57 @@ def main() -> None:
     model.load_state_dict(checkpoint["state_dict"], strict=True)
     model.eval()
     plan_requested = bool(args.plan)
-    plan_applied = plan_requested and trainer_config.policy_type == "reveal_state"
+    plan_applied = plan_requested and planner_enabled(trainer_config, during_eval=True)
     planning_note = None
-    if plan_requested and trainer_config.policy_type != "reveal_state":
+    if plan_requested and not policy_supports_planning(trainer_config.policy_type):
         plan_applied = False
         planning_note = "Planner requested for a backbone-only checkpoint; evaluating the backbone policy only."
-    elif plan_requested and trainer_config.policy_type == "reveal_state" and not args.allow_unsupervised_planning:
-        plan_applied = False
-        planning_note = (
-            "RLBench batches do not provide reveal supervision. Unsupervised reveal planning was disabled; "
-            "use --allow-unsupervised-planning to override."
-        )
-
-    obs_config = create_obs_config(
-        ["front", "wrist_left", "wrist_right"],
-        [args.resolution, args.resolution],
-        "BIMANUAL_PERACT",
-        "bimanual",
-    )
-    action_mode = BimanualMoveArmThenGripper(
-        BimanualEndEffectorPoseViaPlanning(absolute_mode=True, frame="world", collision_checking=False),
-        BimanualDiscrete(),
-    )
-    env = Environment(
-        action_mode=action_mode,
-        obs_config=obs_config,
-        headless=args.headless,
-        robot_setup="dual_panda",
-    )
+    elif plan_requested and trainer_config.planner_mode == "off":
+        planning_note = "Planner requested, but the checkpoint configuration sets planner_mode=off."
+    elif plan_requested and not args.allow_unsupervised_planning and trainer_config.planner_mode == "selfsup":
+        planning_note = "Planner is running in self-supervised mode without direct RLBench planner labels."
+
+    camera_spec = default_three_camera_spec(args.resolution)
 
     results: dict[str, Any] = {
         "checkpoint": str(Path(args.checkpoint).resolve()),
         "plan_requested": plan_requested,
         "plan_applied": plan_applied,
+        "planner_mode": trainer_config.planner_mode,
         "support_mode_conditioning": not args.disable_support_mode_conditioning,
         "episodes_per_task": args.episodes_per_task,
         "episode_length": args.episode_length,
         "resolution": args.resolution,
+        "cameras": list(camera_spec.cameras),
         "tasks": {},
     }
     if planning_note is not None:
         results["planning_note"] = planning_note
 
-    try:
-        env.launch()
-        for task_name in args.tasks:
-            task_class = task_file_to_task_class(task_name, bimanual=True)
+    for task_name in args.tasks:
+        task_successes: list[float] = []
+        task_returns: list[float] = []
+        env: Environment2Robots | None = None
+        try:
+            task_class = task_file_to_task_class(task_name)
+            obs_config = create_obs_config(
+                list(camera_spec.upstream_cameras),
+                [args.resolution, args.resolution],
+                "PERACT_BC",
+            )
+            action_mode = MoveArmThenGripper2Robots(
+                EndEffectorPoseViaPlanning2Robots(absolute_mode=True, frame="world", collision_checking=False),
+                Discrete2Robots(),
+            )
+            env = Environment2Robots(
+                action_mode=action_mode,
+                obs_config=obs_config,
+                headless=args.headless,
+                robot_setup="panda",
+                task_name=task_class.__name__,
+            )
+            env.launch()
             task = env.get_task(task_class)
-            task_successes: list[float] = []
-            task_returns: list[float] = []
             for _ in range(args.episodes_per_task):
                 descriptions, obs = task.reset()
                 language_goal = _episode_language_goal(descriptions)
@@ -157,7 +176,7 @@ def main() -> None:
                             dtype=proprio.dtype,
                         )
                     with torch.no_grad():
-                        if trainer_config.policy_type == "reveal_state":
+                        if policy_supports_planning(trainer_config.policy_type):
                             outputs = model(
                                 images=images,
                                 proprio=proprio,
@@ -181,12 +200,12 @@ def main() -> None:
                     step_action = chosen_chunk[0, 0].detach().float().cpu().numpy()
                     if history_steps > 0:
                         if len(history_images) >= history_steps:
-                            history_images = history_images[-history_steps + 1 :]
-                            history_proprio = history_proprio[-history_steps + 1 :]
+                            keep = max(history_steps - 1, 0)
+                            history_images = history_images[-keep:] if keep > 0 else []
+                            history_proprio = history_proprio[-keep:] if keep > 0 else []
                         history_images.append(images[0].detach().cpu().numpy())
                         history_proprio.append(proprio[0].detach().cpu().numpy())
-                    env_action = absolute_action_from_delta(obs, step_action, ignore_collisions=True)
-                    obs, reward, done = task.step(env_action)
+                    obs, reward, done = _step_bimanual_chunk(task, obs, step_action)
                     total_reward += float(reward)
                     if reward >= 1.0:
                         success = 1.0
@@ -195,13 +214,17 @@ def main() -> None:
                 task_successes.append(success)
                 task_returns.append(total_reward)
             results["tasks"][task_name] = {
+                "task_class": task_class.__name__,
                 "successes": task_successes,
                 "returns": task_returns,
                 "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,
             }
-    finally:
-        env.shutdown()
+        except Exception as exc:
+            results["tasks"][task_name] = {"error": str(exc), "mean_success": 0.0, "mean_return": 0.0}
+        finally:
+            if env is not None:
+                env.shutdown()
 
     task_scores = [task_data["mean_success"] for task_data in results["tasks"].values()]
     results["mean_success"] = float(np.mean(task_scores)) if task_scores else 0.0
@@ -222,7 +245,10 @@ def main() -> None:
         "",
     ]
     for task_name, task_data in results["tasks"].items():
-        lines.append(f"- `{task_name}`: mean_success={task_data['mean_success']:.3f}, returns={task_data['returns']}")
+        if "error" in task_data:
+            lines.append(f"- `{task_name}`: error={task_data['error']}")
+        else:
+            lines.append(f"- `{task_name}`: mean_success={task_data['mean_success']:.3f}, returns={task_data['returns']}")
     (output_dir / "rollout_eval.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
     print(json.dumps(results, indent=2))
 

code/reveal_vla_bimanual/models/__init__.py

@@ -1,10 +1,11 @@
-from models.action_decoder import ACTBimanualChunkDecoder, ChunkDecoderConfig
+from models.action_decoder import ACTBimanualChunkDecoder, ChunkDecoderConfig, InteractionChunkDecoder
 from models.backbones import FrozenVLBackbone, FrozenVLBackboneConfig
 from models.multiview_fusion import MultiViewFusion, MultiViewFusionConfig
-from models.planner import PlannerConfig, RevealPlanner
-from models.policy import BackboneOnlyPolicy, RevealBimanualPolicy
-from models.reveal_head import RevealHeadConfig, RevealStateHead
-from models.world_model import RevealWM, RevealWMConfig
+from models.observation_memory import InteractionObservationMemory, ObservationMemory, ObservationMemoryConfig
+from models.planner import InteractionPlanner, PlannerConfig, RevealPlanner
+from models.policy import BackboneOnlyPolicy, InteractionBimanualPolicy, RevealBimanualPolicy
+from models.reveal_head import InteractionStateHead, RevealHeadConfig, RevealStateHead
+from models.world_model import InteractionRolloutModel, RevealWM, RevealWMConfig
 
 __all__ = [
     "ACTBimanualChunkDecoder",
@@ -12,8 +13,16 @@ __all__ = [
     "ChunkDecoderConfig",
     "FrozenVLBackbone",
     "FrozenVLBackboneConfig",
+    "InteractionBimanualPolicy",
+    "InteractionChunkDecoder",
+    "InteractionObservationMemory",
+    "InteractionPlanner",
+    "InteractionRolloutModel",
+    "InteractionStateHead",
     "MultiViewFusion",
     "MultiViewFusionConfig",
+    "ObservationMemory",
+    "ObservationMemoryConfig",
     "PlannerConfig",
     "RevealBimanualPolicy",
     "RevealHeadConfig",

code/reveal_vla_bimanual/models/action_decoder.py

@@ -17,6 +17,8 @@ class ChunkDecoderConfig:
     action_dim: int = 14
     arm_action_dim: int = 7
     num_candidates: int = 8
+    num_phases: int = 5
+    num_arm_roles: int = 4
 
 
 class ACTBimanualChunkDecoder(nn.Module):
@@ -157,3 +159,225 @@ class ACTBimanualChunkDecoder(nn.Module):
         candidates = action_mean.unsqueeze(1) + noise * std.unsqueeze(1)
         candidates[:, 0] = action_mean
         return candidates
+
+
+class InteractionChunkDecoder(nn.Module):
+    def __init__(self, config: ChunkDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        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.right_decoder = nn.TransformerDecoder(decoder_layer, num_layers=config.num_layers)
+        left_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.left_decoder = nn.TransformerDecoder(left_layer, num_layers=config.num_layers)
+        self.query_embed = nn.Embedding(config.chunk_size, config.hidden_dim)
+        self.proposal_queries = nn.Embedding(config.num_candidates, 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.right_mean = nn.Linear(config.hidden_dim, config.arm_action_dim)
+        self.right_log_std = nn.Linear(config.hidden_dim, config.arm_action_dim)
+        self.left_mean = nn.Linear(config.hidden_dim, config.action_dim - config.arm_action_dim)
+        self.left_log_std = nn.Linear(config.hidden_dim, config.action_dim - config.arm_action_dim)
+        self.proposal_score = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim * 3),
+            nn.Linear(config.hidden_dim * 3, 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,
+    ) -> tuple[Tensor, Tensor, Tensor | None]:
+        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)
+            return zero_phase, zero_roles, None
+        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)
+        phase_context = self.phase_adapter(phase_probs)
+        role_context = self.role_adapter(arm_role_probs)
+        return phase_context, role_context, interaction_state.get("interaction_tokens")
+
+    def _decode_from_queries(
+        self,
+        queries: Tensor,
+        decoder_memory: Tensor,
+        phase_context: Tensor,
+        role_context: Tensor,
+        interaction_context: Tensor,
+    ) -> dict[str, Tensor]:
+        phase_bias = phase_context.unsqueeze(1)
+        right_queries = (
+            queries
+            + phase_bias
+            + role_context[:, 0].unsqueeze(1)
+            + self.arm_identity.weight[0].view(1, 1, -1).to(dtype=queries.dtype)
+        )
+        left_queries = (
+            queries
+            + phase_bias
+            + role_context[:, 1].unsqueeze(1)
+            + self.arm_identity.weight[1].view(1, 1, -1).to(dtype=queries.dtype)
+        )
+        right_tokens = self.right_decoder(right_queries, decoder_memory)
+        left_tokens = self.left_decoder(left_queries, torch.cat([decoder_memory, right_tokens], dim=1))
+        context = interaction_context.unsqueeze(1).expand(-1, queries.shape[1], -1)
+        coordination_input = torch.cat([right_tokens, left_tokens, context], dim=-1)
+        coordination = torch.tanh(self.coordination(coordination_input))
+        right_tokens = right_tokens + coordination
+        left_tokens = left_tokens + coordination
+        action_mean = torch.cat([self.right_mean(right_tokens), self.left_mean(left_tokens)], dim=-1)
+        action_log_std = torch.cat(
+            [self.right_log_std(right_tokens), self.left_log_std(left_tokens)],
+            dim=-1,
+        ).clamp(min=-5.0, max=2.0)
+        pooled_features = torch.cat(
+            [right_tokens.mean(dim=1), left_tokens.mean(dim=1), coordination.mean(dim=1)],
+            dim=-1,
+        )
+        return {
+            "right_tokens": right_tokens,
+            "left_tokens": left_tokens,
+            "coordination_tokens": coordination,
+            "decoded_tokens": torch.cat([right_tokens, left_tokens], dim=-1),
+            "action_mean": action_mean,
+            "action_log_std": action_log_std,
+            "proposal_score": self.proposal_score(pooled_features).squeeze(-1),
+        }
+
+    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,
+    ) -> 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_tokens = self._conditioning(
+            interaction_state=interaction_state,
+            batch_size=batch_size,
+            device=scene_tokens.device,
+            dtype=dtype,
+        )
+
+        decoder_memory = scene_tokens
+        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)
+
+        if interaction_tokens is not None and interaction_tokens.numel() > 0:
+            interaction_context = interaction_tokens.mean(dim=1)
+        elif reveal_tokens is not None and reveal_tokens.numel() > 0:
+            interaction_context = reveal_tokens.mean(dim=1)
+        else:
+            interaction_context = scene_tokens.mean(dim=1)
+        interaction_context = self.context_proj(interaction_context)
+
+        base_queries = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        decoded = self._decode_from_queries(
+            queries=base_queries,
+            decoder_memory=decoder_memory,
+            phase_context=phase_context,
+            role_context=role_context,
+            interaction_context=interaction_context,
+        )
+
+        num_candidates = self.config.num_candidates
+        proposal_bias = self.proposal_queries.weight.view(1, num_candidates, 1, -1).expand(
+            batch_size, -1, self.config.chunk_size, -1
+        )
+        candidate_queries = base_queries.unsqueeze(1) + proposal_bias
+        flat_queries = candidate_queries.reshape(batch_size * num_candidates, self.config.chunk_size, self.config.hidden_dim)
+        flat_memory = decoder_memory.unsqueeze(1).expand(-1, num_candidates, -1, -1).reshape(
+            batch_size * num_candidates, decoder_memory.shape[1], decoder_memory.shape[2]
+        )
+        flat_phase = phase_context.unsqueeze(1).expand(-1, num_candidates, -1).reshape(
+            batch_size * num_candidates, self.config.hidden_dim
+        )
+        flat_roles = role_context.unsqueeze(1).expand(-1, num_candidates, -1, -1).reshape(
+            batch_size * num_candidates, 2, self.config.hidden_dim
+        )
+        flat_context = interaction_context.unsqueeze(1).expand(-1, num_candidates, -1).reshape(
+            batch_size * num_candidates, self.config.hidden_dim
+        )
+        candidate_decoded = self._decode_from_queries(
+            queries=flat_queries,
+            decoder_memory=flat_memory,
+            phase_context=flat_phase,
+            role_context=flat_roles,
+            interaction_context=flat_context,
+        )
+
+        proposal_deltas = candidate_decoded["action_mean"].view(
+            batch_size,
+            num_candidates,
+            self.config.chunk_size,
+            self.config.action_dim,
+        )
+        proposal_logits = candidate_decoded["proposal_score"].view(batch_size, num_candidates)
+        proposal_candidates = decoded["action_mean"].unsqueeze(1) + 0.35 * torch.tanh(proposal_deltas)
+        proposal_candidates[:, 0] = decoded["action_mean"]
+        proposal_logits[:, 0] = decoded["proposal_score"]
+        decoded["proposal_candidates"] = proposal_candidates
+        decoded["proposal_logits"] = proposal_logits
+        return decoded
+
+    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/observation_memory.py

@@ -12,6 +12,9 @@ class ObservationMemoryConfig:
     history_steps: int = 2
     num_layers: int = 1
     dropout: float = 0.1
+    memory_bank_size: int = 4
+    num_heads: int = 4
+    max_history_steps: int = 8
 
 
 class ObservationMemory(nn.Module):
@@ -52,5 +55,86 @@ class ObservationMemory(nn.Module):
             "memory_sequence": memory_sequence,
             "memory_state": final_state,
             "memory_token": self.token_proj(final_state).unsqueeze(1),
+            "memory_tokens": self.token_proj(final_state).unsqueeze(1),
             "memory_uncertainty": torch.nn.functional.softplus(self.uncertainty_head(final_state)).squeeze(-1),
         }
+
+
+class InteractionObservationMemory(nn.Module):
+    def __init__(self, config: ObservationMemoryConfig) -> None:
+        super().__init__()
+        self.config = config
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.hidden_dim,
+            nhead=config.num_heads,
+            dim_feedforward=config.hidden_dim * 4,
+            dropout=config.dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.sequence_encoder = nn.TransformerEncoder(encoder_layer, num_layers=max(1, config.num_layers))
+        self.position_embedding = nn.Parameter(
+            torch.randn(1, config.max_history_steps + 1, config.hidden_dim) * 0.02
+        )
+        self.bank_queries = nn.Parameter(torch.randn(config.memory_bank_size, config.hidden_dim) * 0.02)
+        self.bank_attention = nn.MultiheadAttention(
+            embed_dim=config.hidden_dim,
+            num_heads=config.num_heads,
+            dropout=config.dropout,
+            batch_first=True,
+        )
+        self.bank_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.token_proj = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        self.uncertainty_head = nn.Sequential(
+            nn.LayerNorm(config.hidden_dim),
+            nn.Linear(config.hidden_dim, 1),
+        )
+
+    def _truncate_history(self, history_scene_tokens: Tensor | None) -> Tensor | None:
+        if history_scene_tokens is None or history_scene_tokens.numel() == 0:
+            return history_scene_tokens
+        if history_scene_tokens.shape[1] <= self.config.history_steps:
+            return history_scene_tokens
+        return history_scene_tokens[:, -self.config.history_steps :]
+
+    def forward(
+        self,
+        scene_tokens: Tensor,
+        history_scene_tokens: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        pooled_current = scene_tokens.mean(dim=1, keepdim=True)
+        history_scene_tokens = self._truncate_history(history_scene_tokens)
+        if history_scene_tokens is not None and history_scene_tokens.numel() > 0:
+            history_pooled = history_scene_tokens.mean(dim=2)
+            sequence = torch.cat([history_pooled, pooled_current], dim=1)
+        else:
+            sequence = pooled_current
+
+        seq_len = sequence.shape[1]
+        if seq_len > self.position_embedding.shape[1]:
+            raise ValueError(
+                f"Sequence length {seq_len} exceeds configured maximum {self.position_embedding.shape[1]}"
+            )
+        encoded = self.sequence_encoder(sequence + self.position_embedding[:, :seq_len])
+        batch_size = encoded.shape[0]
+        queries = self.bank_queries.unsqueeze(0).expand(batch_size, -1, -1)
+        bank_tokens, _ = self.bank_attention(queries, encoded, encoded)
+        bank_tokens = bank_tokens + self.bank_mlp(bank_tokens)
+        projected_bank = self.token_proj(bank_tokens)
+        pooled_bank = projected_bank.mean(dim=1)
+        return {
+            "memory_sequence": encoded,
+            "memory_state": encoded[:, -1],
+            "memory_token": pooled_bank.unsqueeze(1),
+            "memory_tokens": projected_bank,
+            "memory_uncertainty": torch.nn.functional.softplus(self.uncertainty_head(pooled_bank)).squeeze(-1),
+        }

code/reveal_vla_bimanual/models/planner.py

@@ -11,6 +11,7 @@ class PlannerConfig:
     hidden_dim: int = 512
     num_candidates: int = 8
     action_dim: int = 14
+    num_support_modes: int = 3
     utility_margin: float = 0.1
     corridor_weight: float = 1.0
     persistence_weight: float = 0.5
@@ -19,6 +20,10 @@ class PlannerConfig:
     disturbance_weight: float = 0.75
     reocclusion_weight: float = 0.5
     visibility_weight: float = 0.25
+    num_heads: int = 4
+    num_layers: int = 2
+    num_phases: int = 5
+    num_arm_roles: int = 4
 
 
 class RevealPlanner(nn.Module):
@@ -87,3 +92,113 @@ class RevealPlanner(nn.Module):
             "best_indices": best_idx,
             "best_chunk": candidate_chunks[batch_indices, best_idx],
         }
+
+
+class InteractionPlanner(nn.Module):
+    def __init__(self, config: PlannerConfig) -> None:
+        super().__init__()
+        self.config = config
+        step_dim = (
+            config.action_dim
+            + config.num_phases
+            + (2 * config.num_arm_roles)
+            + config.num_support_modes
+            + 7
+        )
+        self.step_proj = nn.Sequential(
+            nn.LayerNorm(step_dim),
+            nn.Linear(step_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.hidden_dim,
+            nhead=config.num_heads,
+            dim_feedforward=config.hidden_dim * 4,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.sequence_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_dim) * 0.02)
+        self.success_head = nn.Linear(config.hidden_dim, 1)
+        self.risk_head = nn.Linear(config.hidden_dim, 1)
+        self.score_head = nn.Linear(config.hidden_dim, 1)
+
+    def _mean_field(self, tensor: Tensor) -> Tensor:
+        return tensor.mean(dim=(-1, -2))
+
+    def summarize_trajectory(self, candidate_chunks: Tensor, rollout_state: dict[str, Tensor]) -> Tensor:
+        horizon = min(candidate_chunks.shape[2], rollout_state["phase_logits"].shape[2])
+        candidate_steps = candidate_chunks[:, :, :horizon]
+        phase_probs = rollout_state["phase_logits"][:, :, :horizon].softmax(dim=-1)
+        support_probs = rollout_state["support_mode_logits"][:, :, :horizon].softmax(dim=-1)
+        arm_role_probs = rollout_state["arm_role_logits"][:, :, :horizon].softmax(dim=-1).flatten(start_dim=-2)
+        target_mean = self._mean_field(rollout_state["target_field"][:, :, :horizon].sigmoid())
+        feasibility_mean = self._mean_field(rollout_state["actor_feasibility_field"][:, :, :horizon].sigmoid())
+        persistence_mean = self._mean_field(rollout_state["persistence_field"][:, :, :horizon])
+        risk_mean = self._mean_field(rollout_state["risk_field"][:, :, :horizon])
+        uncertainty_mean = self._mean_field(rollout_state["uncertainty_field"][:, :, :horizon])
+        role_gap = (
+            rollout_state["arm_role_logits"][:, :, :horizon, 0].softmax(dim=-1)
+            - rollout_state["arm_role_logits"][:, :, :horizon, 1].softmax(dim=-1)
+        ).abs().mean(dim=-1, keepdim=True)
+        return torch.cat(
+            [
+                candidate_steps,
+                phase_probs,
+                arm_role_probs,
+                support_probs,
+                target_mean,
+                feasibility_mean,
+                persistence_mean,
+                risk_mean,
+                uncertainty_mean,
+                role_gap,
+            ],
+            dim=-1,
+        )
+
+    def score_rollouts(
+        self,
+        rollout_state: dict[str, Tensor],
+        candidate_chunks: Tensor,
+        proposal_logits: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        features = self.summarize_trajectory(candidate_chunks, rollout_state)
+        batch_size, num_candidates, horizon, _ = features.shape
+        flat_features = features.view(batch_size * num_candidates, horizon, -1)
+        hidden_steps = self.step_proj(flat_features)
+        cls = self.cls_token.expand(batch_size * num_candidates, -1, -1)
+        encoded = self.sequence_encoder(torch.cat([cls, hidden_steps], dim=1))
+        pooled = encoded[:, 0]
+        success_logits = self.success_head(pooled).view(batch_size, num_candidates).squeeze(-1)
+        risk_values = torch.sigmoid(self.risk_head(pooled)).view(batch_size, num_candidates).squeeze(-1)
+        utility_scores = self.score_head(pooled).view(batch_size, num_candidates).squeeze(-1)
+        utility_scores = utility_scores + success_logits.sigmoid() - risk_values
+        if proposal_logits is not None and proposal_logits.shape == utility_scores.shape:
+            utility_scores = utility_scores + self.config.proposal_weight * proposal_logits.sigmoid()
+        return {
+            "planner_features": features.mean(dim=2),
+            "planner_hidden": pooled.view(batch_size, num_candidates, -1),
+            "success_logits": success_logits,
+            "risk_values": risk_values,
+            "utility_scores": utility_scores,
+        }
+
+    def select_best(
+        self,
+        candidate_chunks: Tensor,
+        rollout_state: dict[str, Tensor],
+        proposal_logits: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        outputs = self.score_rollouts(
+            rollout_state=rollout_state,
+            candidate_chunks=candidate_chunks,
+            proposal_logits=proposal_logits,
+        )
+        best_idx = outputs["utility_scores"].argmax(dim=-1)
+        batch_indices = torch.arange(candidate_chunks.shape[0], device=candidate_chunks.device)
+        return {
+            **outputs,
+            "best_indices": best_idx,
+            "best_chunk": candidate_chunks[batch_indices, best_idx],
+        }

code/reveal_vla_bimanual/models/policy.py

@@ -6,13 +6,13 @@ from typing import Sequence
 import torch
 from torch import Tensor, nn
 
-from models.action_decoder import ACTBimanualChunkDecoder, ChunkDecoderConfig
+from models.action_decoder import ACTBimanualChunkDecoder, ChunkDecoderConfig, InteractionChunkDecoder
 from models.backbones import FrozenVLBackbone, FrozenVLBackboneConfig
 from models.multiview_fusion import MultiViewFusion, MultiViewFusionConfig
-from models.observation_memory import ObservationMemory, ObservationMemoryConfig
-from models.planner import PlannerConfig, RevealPlanner
-from models.reveal_head import RevealHeadConfig, RevealStateHead
-from models.world_model import RevealWM, RevealWMConfig
+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
 
 
 @dataclass
@@ -204,3 +204,138 @@ class RevealBimanualPolicy(BackboneOnlyPolicy):
             outputs["planner_scores"] = selected["utility_scores"]
             outputs["best_candidate_indices"] = selected["best_indices"]
         return outputs
+
+
+class InteractionBimanualPolicy(BackboneOnlyPolicy):
+    def __init__(self, config: PolicyConfig) -> None:
+        super().__init__(config)
+        self.memory = InteractionObservationMemory(config.memory)
+        self.decoder = InteractionChunkDecoder(config.decoder)
+        self.interaction_head = InteractionStateHead(config.reveal_head)
+        self.world_model = InteractionRolloutModel(config.world_model)
+        self.planner = InteractionPlanner(config.planner)
+
+    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]:
+        return {key: self._tile_tensor(value, num_candidates) for key, value in state.items()}
+
+    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,
+        plan: bool = True,
+        support_mode_conditioning: bool = True,
+        candidate_chunks_override: Tensor | None = None,
+        use_interaction_head: bool = True,
+        use_role_tokens: bool = True,
+        history_steps_override: int | None = None,
+    ) -> dict[str, Tensor]:
+        scene_tokens = self.encode_scene(images, proprio, texts=texts, language_tokens=language_tokens)
+        history_scene_tokens = self.encode_history(
+            history_images,
+            history_proprio,
+            texts=texts,
+            language_tokens=language_tokens,
+        )
+        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:]
+        memory_output = self.memory(scene_tokens, history_scene_tokens=history_scene_tokens)
+
+        interaction_state = None
+        if use_interaction_head:
+            interaction_state = self.interaction_head(
+                scene_tokens,
+                memory_tokens=memory_output["memory_tokens"],
+            )
+            interaction_state["memory_tokens"] = memory_output["memory_tokens"]
+            interaction_state["memory_token"] = memory_output["memory_token"]
+
+        if interaction_state is not None and not use_role_tokens:
+            interaction_state = dict(interaction_state)
+            interaction_state["arm_role_logits"] = torch.zeros_like(interaction_state["arm_role_logits"])
+
+        decoded = self.decoder(
+            scene_tokens,
+            interaction_state=interaction_state,
+            memory_tokens=memory_output["memory_tokens"],
+        )
+        outputs = {
+            **decoded,
+            "scene_tokens": scene_tokens,
+            "history_scene_tokens": history_scene_tokens,
+            "memory_output": memory_output,
+            "memory_uncertainty": memory_output["memory_uncertainty"],
+            "interaction_state": interaction_state,
+            "reveal_state": interaction_state,
+        }
+
+        if plan:
+            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 interaction_state is None:
+                outputs["planned_chunk"] = outputs["action_mean"]
+                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,
+                )
+                outputs["planned_rollout"] = {}
+                return outputs
+
+            batch_size, num_candidates, chunk_size, action_dim = candidate_chunks.shape
+            flat_chunks = candidate_chunks.view(batch_size * num_candidates, chunk_size, action_dim)
+            tiled_scene = self._tile_tensor(scene_tokens, num_candidates)
+            planning_state = interaction_state
+            if not support_mode_conditioning:
+                planning_state = dict(interaction_state)
+                planning_state["support_mode_logits"] = torch.zeros_like(interaction_state["support_mode_logits"])
+            tiled_state = self._tile_state(planning_state, num_candidates)
+            tiled_memory_tokens = self._tile_tensor(memory_output["memory_tokens"], num_candidates)
+            rollout = self.world_model(
+                scene_tokens=tiled_scene,
+                interaction_state=tiled_state,
+                action_chunk=flat_chunks,
+                memory_tokens=tiled_memory_tokens,
+            )
+            reshaped_rollout = {
+                key: value.view(batch_size, num_candidates, *value.shape[1:]) for key, value in rollout.items()
+            }
+            selected = self.planner.select_best(
+                candidate_chunks=candidate_chunks,
+                rollout_state=reshaped_rollout,
+                proposal_logits=proposal_logits,
+            )
+            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_scores"]
+            outputs["best_candidate_indices"] = selected["best_indices"]
+        return outputs

code/reveal_vla_bimanual/models/reveal_head.py

@@ -17,6 +17,9 @@ class RevealHeadConfig:
     field_size: int = 16
     num_heads: int = 4
     predict_belief_map: bool = False
+    num_phases: int = 5
+    num_arm_roles: int = 4
+    num_interaction_tokens: int = 8
 
 
 class RevealStateHead(nn.Module):
@@ -116,3 +119,201 @@ class RevealStateHead(nn.Module):
         if self.config.predict_belief_map:
             output["belief_map"] = belief_map
         return output
+
+
+class InteractionFieldDecoder(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.target_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.actor_feasibility_field = nn.Conv2d(config.hidden_dim, 2, kernel_size=1)
+        self.persistence_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.risk_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.uncertainty_field = nn.Conv2d(config.hidden_dim, 1, kernel_size=1)
+        self.compat_access_field = nn.Conv2d(config.hidden_dim, config.num_support_modes, kernel_size=1)
+        self.compat_persistence = nn.Conv2d(config.hidden_dim, config.num_support_modes, 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 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)
+        summary = self.summary_proj(summary_input)
+
+        target_field = self.target_field(grid)
+        actor_feasibility_field = self.actor_feasibility_field(grid)
+        persistence_field = torch.sigmoid(self.persistence_field(grid))
+        risk_field = torch.sigmoid(self.risk_field(grid))
+        uncertainty_field = F.softplus(self.uncertainty_field(grid))
+
+        access_field = self.compat_access_field(grid)
+        corridor_source = access_field.amax(dim=-2)
+        corridor_logits = F.interpolate(
+            corridor_source,
+            size=self.config.num_approach_templates,
+            mode="linear",
+            align_corners=False,
+        )
+        compatibility_persistence = torch.sigmoid(self.compat_persistence(grid))
+        access_prob = torch.sigmoid(access_field)
+        weighted_persistence = (compatibility_persistence * 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 = risk_field.mean(dim=(-1, -2)).squeeze(1)
+        belief_map = target_field
+        if belief_map.shape[-1] != self.config.belief_map_size:
+            belief_map = F.interpolate(
+                belief_map,
+                size=(self.config.belief_map_size, self.config.belief_map_size),
+                mode="bilinear",
+                align_corners=False,
+            )
+
+        arm_identity = self.arm_identity.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        if interaction_tokens.shape[1] >= 2:
+            arm_tokens = interaction_tokens[:, :2] + arm_identity
+        else:
+            arm_tokens = pooled_interaction.unsqueeze(1).expand(-1, 2, -1) + arm_identity
+        arm_role_input = torch.cat(
+            [arm_tokens, summary.unsqueeze(1).expand(-1, arm_tokens.shape[1], -1)],
+            dim=-1,
+        )
+        arm_role_logits = self.arm_role_head(arm_role_input)
+        reocclusion_logit = self.reocclusion_head(summary_input)
+
+        output = {
+            "phase_logits": self.phase_head(summary_input),
+            "arm_role_logits": arm_role_logits,
+            "target_field": target_field,
+            "actor_feasibility_field": actor_feasibility_field,
+            "persistence_field": persistence_field,
+            "risk_field": risk_field,
+            "uncertainty_field": uncertainty_field,
+            "interaction_tokens": interaction_tokens,
+            "field_tokens": field_tokens,
+            "latent_summary": summary,
+            "support_mode_logits": self.support_mode(summary_input),
+            "corridor_logits": corridor_logits,
+            "persistence_horizon": persistence_horizon,
+            "disturbance_cost": disturbance_cost,
+            "belief_map": belief_map,
+            "reocclusion_logit": reocclusion_logit,
+            "persistence_uncertainty": uncertainty_field.mean(dim=(-1, -2)).squeeze(1),
+            "access_field": access_field,
+            "disturbance_field": risk_field,
+            "uncertainty": uncertainty_field.mean(dim=(-1, -2)).squeeze(1),
+        }
+        if not self.config.predict_belief_map:
+            output.pop("belief_map")
+        return output
+
+
+class InteractionStateHead(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 = InteractionFieldDecoder(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

@@ -5,6 +5,8 @@ from dataclasses import dataclass
 import torch
 from torch import Tensor, nn
 
+from models.reveal_head import InteractionFieldDecoder
+
 
 @dataclass
 class RevealWMConfig:
@@ -13,6 +15,13 @@ class RevealWMConfig:
     num_support_modes: int = 3
     num_approach_templates: int = 32
     rollout_horizon: int = 5
+    field_size: int = 16
+    num_heads: int = 4
+    num_phases: int = 5
+    num_arm_roles: int = 4
+    num_interaction_tokens: int = 8
+    belief_map_size: int = 32
+    predict_belief_map: bool = True
 
 
 class RevealWM(nn.Module):
@@ -78,3 +87,68 @@ class RevealWM(nn.Module):
             "reocclusion_logit": self.reocclusion(rollout),
             "uncertainty": torch.nn.functional.softplus(self.uncertainty(rollout)).squeeze(-1),
         }
+
+
+class InteractionRolloutModel(nn.Module):
+    def __init__(self, config: RevealWMConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.action_encoder = nn.Sequential(
+            nn.LayerNorm(config.action_dim),
+            nn.Linear(config.action_dim, config.hidden_dim),
+            nn.GELU(),
+        )
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.hidden_dim,
+            nhead=config.num_heads,
+            dim_feedforward=config.hidden_dim * 4,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.transition = nn.TransformerEncoder(encoder_layer, num_layers=2)
+        self.token_update = 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 = InteractionFieldDecoder(config)
+
+    def forward(
+        self,
+        scene_tokens: Tensor,
+        interaction_state: dict[str, Tensor],
+        action_chunk: Tensor,
+        memory_tokens: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        if memory_tokens is None:
+            memory_tokens = interaction_state.get("memory_tokens")
+        if memory_tokens is None:
+            memory_tokens = interaction_state.get("memory_token")
+        current_tokens = interaction_state["interaction_tokens"]
+        outputs: dict[str, list[Tensor]] = {}
+
+        for step in range(action_chunk.shape[1]):
+            action_token = self.action_encoder(action_chunk[:, step]).unsqueeze(1)
+            transition_tokens = current_tokens
+            if memory_tokens is not None:
+                transition_tokens = torch.cat([transition_tokens, memory_tokens], dim=1)
+            transition_tokens = torch.cat([transition_tokens, action_token], dim=1)
+            transitioned = self.transition(transition_tokens)
+            current_tokens = current_tokens + self.token_update(transitioned[:, : current_tokens.shape[1]])
+            decoded = self.decoder(
+                interaction_tokens=current_tokens,
+                scene_tokens=scene_tokens,
+                memory_tokens=memory_tokens,
+            )
+            decoded["memory_token"] = (
+                memory_tokens.mean(dim=1, keepdim=True) if memory_tokens is not None else current_tokens.mean(dim=1, keepdim=True)
+            )
+            decoded["memory_tokens"] = memory_tokens if memory_tokens is not None else current_tokens[:, :1]
+            for key, value in decoded.items():
+                outputs.setdefault(key, []).append(value)
+
+        stacked: dict[str, Tensor] = {}
+        for key, values in outputs.items():
+            stacked[key] = torch.stack(values, dim=1)
+        return stacked

code/reveal_vla_bimanual/scripts/setup_env_a_rlbench.sh

@@ -57,11 +57,28 @@ run_in_env python -m pip install -U pip setuptools wheel
 run_in_env python -m pip install --force-reinstall --no-deps numpy==1.26.4 Pillow==12.1.1
 run_in_env python -m pip install --no-deps moviepy==2.2.1 timeout-decorator==0.5.0 opencv-python==4.10.0.84 pyquaternion==0.9.9 click-prompt==0.5.1
 run_in_env python -m pip install --no-deps poetry-core
+run_in_env python -m pip install gymnasium==1.0.0a2 open3d==0.19.0 segment-anything==1.0 transforms3d==0.4.1 openai==0.28.1
+PERACT_ROOT="${ROOT_DIR}/third_party/peract_bimanual"
+YARR_ROOT="${ROOT_DIR}/third_party/YARR"
+PYREP_ROOT="${ROOT_DIR}/third_party/PyRep"
+RLBENCH_ROOT="${ROOT_DIR}/third_party/RLBench"
+if [[ -f "${PERACT_ROOT}/YARR/setup.py" ]]; then
+  YARR_ROOT="${PERACT_ROOT}/YARR"
+fi
+if [[ -f "${PERACT_ROOT}/PyRep/setup.py" ]]; then
+  PYREP_ROOT="${PERACT_ROOT}/PyRep"
+fi
+if [[ -f "${PERACT_ROOT}/RLBench/setup.py" ]]; then
+  RLBENCH_ROOT="${PERACT_ROOT}/RLBench"
+fi
+if [[ ! -f "${PERACT_ROOT}/pyproject.toml" && ! -f "${PERACT_ROOT}/setup.py" && -f "${PERACT_ROOT}/peract/setup.py" ]]; then
+  PERACT_ROOT="${PERACT_ROOT}/peract"
+fi
 run_in_env python -m pip install --no-build-isolation -e "${PROJECT_DIR}"
-run_in_env python -m pip install --no-build-isolation --no-deps -e "${ROOT_DIR}/third_party/YARR"
-run_in_env python -m pip install --no-build-isolation --no-deps -e "${ROOT_DIR}/third_party/PyRep"
-run_in_env python -m pip install --no-build-isolation --no-deps -e "${ROOT_DIR}/third_party/RLBench"
-run_in_env python -m pip install --no-build-isolation --no-deps -e "${ROOT_DIR}/third_party/peract_bimanual"
+run_in_env python -m pip install --no-build-isolation --no-deps -e "${YARR_ROOT}"
+run_in_env python -m pip install --no-build-isolation --no-deps -e "${PYREP_ROOT}"
+run_in_env python -m pip install --no-build-isolation --no-deps -e "${RLBENCH_ROOT}"
+run_in_env python -m pip install --no-build-isolation --no-deps -e "${PERACT_ROOT}"
 
 printf 'RLBench env ready at %s\n' "${ENV_PREFIX}"
 printf 'Activate with:\n'

code/reveal_vla_bimanual/scripts/setup_rlbench_headless_x.sh

@@ -9,9 +9,19 @@ export DEBIAN_FRONTEND=noninteractive
 
 apt-get update
 apt-get install -y \
+  libxcb-cursor0 \
+  libxcb-icccm4 \
+  libxcb-image0 \
+  libxcb-keysyms1 \
+  libxcb-randr0 \
+  libxcb-render-util0 \
+  libxcb-xinerama0 \
+  libxrender1 \
   libxkbcommon0 \
   libxkbcommon-x11-0 \
   mesa-utils \
+  xvfb \
+  x11-utils \
   x11-xserver-utils \
   xauth \
   xserver-xorg \

code/reveal_vla_bimanual/sim_rlbench/camera_spec.py

@@ -2,6 +2,21 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+CANONICAL_TO_UPSTREAM_CAMERA = {
+    "front": "front",
+    "wrist_left": "wrist2",
+    "wrist_right": "wrist",
+}
+UPSTREAM_TO_CANONICAL_CAMERA = {value: key for key, value in CANONICAL_TO_UPSTREAM_CAMERA.items()}
+
+
+def canonical_to_upstream_camera(camera_name: str) -> str:
+    return CANONICAL_TO_UPSTREAM_CAMERA.get(camera_name, camera_name)
+
+
+def upstream_to_canonical_camera(camera_name: str) -> str:
+    return UPSTREAM_TO_CANONICAL_CAMERA.get(camera_name, camera_name)
+
 
 @dataclass(frozen=True)
 class RLBenchThreeCameraSpec:
@@ -16,8 +31,12 @@ class RLBenchThreeCameraSpec:
     def wrist_cameras(self) -> tuple[str, str]:
         return self.cameras[1], self.cameras[2]
 
+    @property
+    def upstream_cameras(self) -> tuple[str, str, str]:
+        return tuple(CANONICAL_TO_UPSTREAM_CAMERA.get(camera, camera) for camera in self.cameras)  # type: ignore[return-value]
+
     def hydra_overrides(self, prefix: str = "rlbench") -> list[str]:
-        camera_list = ",".join(self.cameras)
+        camera_list = ",".join(self.upstream_cameras)
         height, width = self.resolution
         return [
             f"{prefix}.cameras=[{camera_list}]",

code/reveal_vla_bimanual/sim_rlbench/dataset.py

@@ -10,10 +10,66 @@ import torch
 from PIL import Image
 from torch.utils.data import Dataset
 
+from sim_rlbench.camera_spec import canonical_to_upstream_camera
+
 
 THREE_CAMERAS: tuple[str, str, str] = ("front", "wrist_left", "wrist_right")
 
 
+def _camera_value(obs: Any, camera_name: str, suffix: str) -> Any:
+    upstream_name = canonical_to_upstream_camera(camera_name)
+    candidate_keys = [
+        f"{upstream_name}_{suffix}",
+        f"{camera_name}_{suffix}",
+    ]
+    if suffix == "point_cloud":
+        candidate_keys.extend(
+            [
+                f"{upstream_name}_pointcloud",
+                f"{camera_name}_pointcloud",
+            ]
+        )
+    for key in candidate_keys:
+        if hasattr(obs, key):
+            return getattr(obs, key)
+    perception_data = getattr(obs, "perception_data", None)
+    if isinstance(perception_data, dict):
+        for key in candidate_keys:
+            if key in perception_data:
+                return perception_data[key]
+    return None
+
+
+def _arm_pose(obs: Any, arm_name: str) -> np.ndarray:
+    key = f"gripper_{arm_name}_pose"
+    if hasattr(obs, key):
+        return np.asarray(getattr(obs, key), dtype=np.float32)
+    arm = getattr(obs, arm_name, None)
+    if arm is not None and hasattr(arm, "gripper_pose"):
+        return np.asarray(arm.gripper_pose, dtype=np.float32)
+    raise AttributeError(f"Observation does not expose pose for arm '{arm_name}'")
+
+
+def _arm_joint_positions(obs: Any, arm_name: str) -> np.ndarray:
+    key = f"joint_positions_{arm_name}"
+    if hasattr(obs, key):
+        return np.asarray(getattr(obs, key), dtype=np.float32)
+    arm = getattr(obs, arm_name, None)
+    if arm is not None and hasattr(arm, "joint_positions"):
+        return np.asarray(arm.joint_positions, dtype=np.float32)
+    raise AttributeError(f"Observation does not expose joint positions for arm '{arm_name}'")
+
+
+def _arm_gripper_open(obs: Any, arm_name: str) -> float:
+    key = f"gripper_{arm_name}_open"
+    if hasattr(obs, key):
+        return float(getattr(obs, key))
+    arm = getattr(obs, arm_name, None)
+    if arm is not None and hasattr(arm, "gripper_open"):
+        return float(arm.gripper_open)
+    raise AttributeError(f"Observation does not expose gripper state for arm '{arm_name}'")
+
+
 def _normalize_quaternion_wxyz(quaternion: np.ndarray) -> np.ndarray:
     quaternion = np.asarray(quaternion, dtype=np.float32)
     return quaternion / max(float(np.linalg.norm(quaternion)), 1e-8)
@@ -90,11 +146,11 @@ def bimanual_proprio_from_obs(
     )
     base = np.concatenate(
         [
-            np.asarray(obs.right.gripper_pose, dtype=np.float32),
-            np.asarray(obs.left.gripper_pose, dtype=np.float32),
-            np.asarray(obs.right.joint_positions, dtype=np.float32),
-            np.asarray(obs.left.joint_positions, dtype=np.float32),
-            np.array([float(obs.right.gripper_open), float(obs.left.gripper_open)], dtype=np.float32),
+            _arm_pose(obs, "right"),
+            _arm_pose(obs, "left"),
+            _arm_joint_positions(obs, "right"),
+            _arm_joint_positions(obs, "left"),
+            np.array([_arm_gripper_open(obs, "right"), _arm_gripper_open(obs, "left")], dtype=np.float32),
             time_feature,
         ],
         axis=0,
@@ -108,10 +164,8 @@ def bimanual_proprio_from_obs(
 def delta_action_from_transition(current_obs: Any, next_obs: Any) -> np.ndarray:
     action_parts: list[np.ndarray] = []
     for arm_name in ("right", "left"):
-        current_arm = getattr(current_obs, arm_name)
-        next_arm = getattr(next_obs, arm_name)
-        current_pose = np.asarray(current_arm.gripper_pose, dtype=np.float32)
-        next_pose = np.asarray(next_arm.gripper_pose, dtype=np.float32)
+        current_pose = _arm_pose(current_obs, arm_name)
+        next_pose = _arm_pose(next_obs, arm_name)
         position_delta = next_pose[:3] - current_pose[:3]
         current_quat = _xyzw_to_wxyz(current_pose[3:])
         next_quat = _xyzw_to_wxyz(next_pose[3:])
@@ -122,7 +176,7 @@ def delta_action_from_transition(current_obs: Any, next_obs: Any) -> np.ndarray:
                 [
                     position_delta.astype(np.float32),
                     delta_rotvec.astype(np.float32),
-                    np.array([float(next_arm.gripper_open)], dtype=np.float32),
+                    np.array([_arm_gripper_open(next_obs, arm_name)], dtype=np.float32),
                 ],
                 axis=0,
             )
@@ -130,39 +184,55 @@ def delta_action_from_transition(current_obs: Any, next_obs: Any) -> np.ndarray:
     return np.concatenate(action_parts, axis=0).astype(np.float32)
 
 
-def absolute_action_from_delta(current_obs: Any, delta_action: Sequence[float], ignore_collisions: bool = True) -> np.ndarray:
+def single_arm_absolute_action_from_delta(
+    current_obs: Any,
+    delta_action: Sequence[float],
+    arm_name: str,
+    ignore_collisions: bool = True,
+) -> np.ndarray:
     delta_action = np.asarray(delta_action, dtype=np.float32)
     if delta_action.shape != (14,):
         raise ValueError(f"Expected delta action shape (14,), received {delta_action.shape}")
+    arm_index = {"right": 0, "left": 1}[arm_name]
+    current_pose = _arm_pose(current_obs, arm_name)
+    offset = arm_index * 7
+    delta_position = delta_action[offset : offset + 3]
+    delta_rotvec = delta_action[offset + 3 : offset + 6]
+    gripper = float(delta_action[offset + 6] > 0.5)
+    current_quat = _xyzw_to_wxyz(current_pose[3:])
+    delta_quat = _rotvec_to_quat_wxyz(delta_rotvec)
+    next_quat = _quat_multiply_wxyz(delta_quat, current_quat)
+    next_pose = np.concatenate(
+        [
+            current_pose[:3] + delta_position,
+            _wxyz_to_xyzw(next_quat),
+        ],
+        axis=0,
+    )
+    return np.concatenate(
+        [
+            next_pose.astype(np.float32),
+            np.array([gripper, float(ignore_collisions)], dtype=np.float32),
+        ],
+        axis=0,
+    )
 
-    env_action: list[float] = []
-    for arm_index, arm_name in enumerate(("right", "left")):
-        arm = getattr(current_obs, arm_name)
-        current_pose = np.asarray(arm.gripper_pose, dtype=np.float32)
-        offset = arm_index * 7
-        delta_position = delta_action[offset : offset + 3]
-        delta_rotvec = delta_action[offset + 3 : offset + 6]
-        gripper = float(delta_action[offset + 6] > 0.5)
-        current_quat = _xyzw_to_wxyz(current_pose[3:])
-        delta_quat = _rotvec_to_quat_wxyz(delta_rotvec)
-        next_quat = _quat_multiply_wxyz(delta_quat, current_quat)
-        next_pose = np.concatenate(
-            [
-                current_pose[:3] + delta_position,
-                _wxyz_to_xyzw(next_quat),
-            ],
-            axis=0,
-        )
-        env_action.extend(next_pose.tolist())
-        env_action.append(gripper)
-        env_action.append(float(ignore_collisions))
-    return np.asarray(env_action, dtype=np.float32)
+
+def absolute_action_from_delta(current_obs: Any, delta_action: Sequence[float], ignore_collisions: bool = True) -> np.ndarray:
+    arm_actions = [
+        single_arm_absolute_action_from_delta(current_obs, delta_action, arm_name, ignore_collisions=ignore_collisions)
+        for arm_name in ("right", "left")
+    ]
+    return np.concatenate(arm_actions, axis=0).astype(np.float32)
 
 
 def stack_live_rgb_obs(obs: Any, cameras: Sequence[str] = THREE_CAMERAS, resolution: int = 224) -> torch.Tensor:
     images: list[np.ndarray] = []
     for camera_name in cameras:
-        rgb = np.asarray(obs.perception_data[f"{camera_name}_rgb"], dtype=np.uint8)
+        rgb_value = _camera_value(obs, camera_name, "rgb")
+        if rgb_value is None:
+            raise KeyError(f"Observation does not expose RGB for camera '{camera_name}'")
+        rgb = np.asarray(rgb_value, dtype=np.uint8)
         image = Image.fromarray(rgb)
         if image.size != (resolution, resolution):
             image = image.resize((resolution, resolution), Image.Resampling.BILINEAR)
@@ -259,6 +329,9 @@ class RLBenchOfflineChunkDataset(Dataset[dict[str, Any]]):
         frames: list[np.ndarray] = []
         for camera_name in self.cameras:
             image_path = episode_dir / f"{camera_name}_rgb" / f"rgb_{step_index:04d}.png"
+            if not image_path.exists():
+                upstream_camera_name = canonical_to_upstream_camera(camera_name)
+                image_path = episode_dir / f"{upstream_camera_name}_rgb" / f"rgb_{step_index:04d}.png"
             image = Image.open(image_path).convert("RGB")
             if image.size != (self.resolution, self.resolution):
                 image = image.resize((self.resolution, self.resolution), Image.Resampling.BILINEAR)

code/reveal_vla_bimanual/sim_rlbench/generate_smoke_dataset.py

@@ -1,17 +1,17 @@
 from __future__ import annotations
 
 import argparse
+import copy
 import pickle
 from pathlib import Path
 
 import numpy as np
 from PIL import Image
-from pyrep.const import RenderMode
-from rlbench.action_modes.action_mode import BimanualMoveArmThenGripper
-from rlbench.action_modes.arm_action_modes import BimanualJointPosition
-from rlbench.action_modes.gripper_action_modes import BimanualDiscrete
+from helpers.utils import create_obs_config
+from rlbench.action_modes.action_mode import MoveArmThenGripper2Robots
+from rlbench.action_modes.arm_action_modes import EndEffectorPoseViaPlanning2Robots
+from rlbench.action_modes.gripper_action_modes import Discrete2Robots
 from rlbench.backend.const import (
-    DEPTH_SCALE,
     EPISODE_FOLDER,
     EPISODES_FOLDER,
     LOW_DIM_PICKLE,
@@ -19,11 +19,24 @@ from rlbench.backend.const import (
     VARIATION_NUMBER,
     VARIATIONS_ALL_FOLDER,
 )
-from rlbench.backend.utils import float_array_to_rgb_image, task_file_to_task_class
-from rlbench.environment import Environment
-from rlbench.observation_config import CameraConfig, ObservationConfig
+from rlbench.backend.utils import task_file_to_task_class
+from rlbench.environments_two_robots import Environment2Robots
 
-from sim_rlbench.camera_spec import default_three_camera_spec
+from sim_rlbench.camera_spec import canonical_to_upstream_camera, default_three_camera_spec
+
+
+def _camera_payload(obs: object, camera_name: str, suffix: str):
+    upstream_name = canonical_to_upstream_camera(camera_name)
+    for key in (f"{upstream_name}_{suffix}", f"{camera_name}_{suffix}"):
+        if hasattr(obs, key):
+            return getattr(obs, key)
+    return None
+
+
+def _scripted_demo(task_env: object, steps_per_episode: int) -> tuple[list[str], list[object]]:
+    descriptions, obs = task_env.reset()
+    demo = [copy.deepcopy(obs) for _ in range(max(steps_per_episode, 2))]
+    return descriptions, demo
 
 
 def _save_demo(demo, episode_path: Path, cameras: list[str]) -> None:
@@ -33,19 +46,14 @@ def _save_demo(demo, episode_path: Path, cameras: list[str]) -> None:
             for dtype in data_types:
                 output_dir = episode_path / f"{camera_name}_{dtype}"
                 output_dir.mkdir(parents=True, exist_ok=True)
-                payload = obs.perception_data.get(f"{camera_name}_{dtype}")
+                payload = _camera_payload(obs, camera_name, dtype)
                 if payload is None:
                     continue
                 if dtype == "rgb":
-                    image = Image.fromarray(payload)
-                elif dtype == "depth":
-                    image = float_array_to_rgb_image(payload, scale_factor=DEPTH_SCALE)
-                elif dtype == "mask":
-                    image = Image.fromarray((payload * 255).astype(np.uint8))
+                    image = Image.fromarray(np.asarray(payload, dtype=np.uint8))
                 else:
-                    raise ValueError(dtype)
+                    image = Image.fromarray((np.asarray(payload) * 255).astype(np.uint8))
                 image.save(output_dir / f"{dtype}_{obs_idx:04d}.png")
-        obs.perception_data.clear()
 
     with (episode_path / LOW_DIM_PICKLE).open("wb") as handle:
         pickle.dump(demo, handle)
@@ -53,43 +61,31 @@ def _save_demo(demo, episode_path: Path, cameras: list[str]) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--task", default="bimanual_lift_ball")
+    parser.add_argument("--task", default="open_drawer")
     parser.add_argument("--episodes", type=int, default=1)
     parser.add_argument("--resolution", type=int, default=224)
-    parser.add_argument("--output-root", default="/workspace/data/rlbench2_smoke")
+    parser.add_argument("--output-root", default="/workspace/data/rlbench_smoke")
+    parser.add_argument("--steps-per-episode", type=int, default=6)
+    parser.add_argument("--try-live-demos", action="store_true")
     args = parser.parse_args()
 
     spec = default_three_camera_spec(args.resolution)
-    camera_config = CameraConfig(
-        rgb=True,
-        depth=True,
-        point_cloud=False,
-        mask=True,
-        image_size=list(spec.resolution),
-        render_mode=RenderMode.OPENGL,
-        masks_as_one_channel=False,
-        depth_in_meters=False,
-    )
-    obs_config = ObservationConfig(
-        camera_configs={camera_name: camera_config for camera_name in spec.cameras},
-        joint_forces=False,
-        joint_positions=True,
-        joint_velocities=True,
-        task_low_dim_state=False,
-        gripper_touch_forces=False,
-        gripper_pose=True,
-        gripper_open=True,
-        gripper_matrix=True,
-        gripper_joint_positions=True,
-        robot_name="bimanual",
+    obs_config = create_obs_config(
+        list(spec.upstream_cameras),
+        [args.resolution, args.resolution],
+        "PERACT_BC",
     )
 
-    task_class = task_file_to_task_class(args.task, bimanual=True)
-    env = Environment(
-        action_mode=BimanualMoveArmThenGripper(BimanualJointPosition(), BimanualDiscrete()),
+    task_class = task_file_to_task_class(args.task)
+    env = Environment2Robots(
+        action_mode=MoveArmThenGripper2Robots(
+            EndEffectorPoseViaPlanning2Robots(absolute_mode=True, frame="world", collision_checking=False),
+            Discrete2Robots(),
+        ),
         obs_config=obs_config,
-        robot_setup="dual_panda",
         headless=True,
+        robot_setup="panda",
+        task_name=task_class.__name__,
     )
     output_root = Path(args.output_root)
     episodes_root = output_root / args.task / VARIATIONS_ALL_FOLDER / EPISODES_FOLDER
@@ -105,8 +101,17 @@ def main() -> None:
             task_env = env.get_task(task_class)
             variation = int(rng.integers(variation_count))
             task_env.set_variation(variation)
-            descriptions, _ = task_env.reset()
-            (demo,) = task_env.get_demos(amount=1, live_demos=True)
+            if args.try_live_demos:
+                collection_mode = "live_demo"
+                try:
+                    (demo,) = task_env.get_demos(amount=1, live_demos=True, max_attempts=1)
+                    descriptions = task_env.get_task_descriptions()
+                except Exception as exc:
+                    collection_mode = f"scripted_fallback: {exc}"
+                    descriptions, demo = _scripted_demo(task_env, steps_per_episode=args.steps_per_episode)
+            else:
+                collection_mode = "scripted"
+                descriptions, demo = _scripted_demo(task_env, steps_per_episode=args.steps_per_episode)
             episode_path = episodes_root / (EPISODE_FOLDER % episode_idx)
             episode_path.mkdir(parents=True, exist_ok=True)
             _save_demo(demo, episode_path, list(spec.cameras))
@@ -115,7 +120,7 @@ def main() -> None:
             with (episode_path / VARIATION_DESCRIPTIONS).open("wb") as handle:
                 pickle.dump(descriptions, handle)
             print(
-                f"[done] wrote {args.task} episode {episode_idx} variation {variation} to {episode_path}",
+                f"[done] wrote {args.task} episode {episode_idx} variation {variation} via {collection_mode} to {episode_path}",
                 flush=True,
             )
     finally:

code/reveal_vla_bimanual/sim_rlbench/obs_adapter.py

@@ -5,7 +5,7 @@ from typing import Any
 
 import numpy as np
 
-from sim_rlbench.camera_spec import RLBenchThreeCameraSpec
+from sim_rlbench.camera_spec import RLBenchThreeCameraSpec, canonical_to_upstream_camera
 
 
 @dataclass
@@ -31,18 +31,66 @@ class CanonicalBimanualObservation:
 
 
 def _camera_rgb(obs: Any, camera_name: str) -> np.ndarray:
-    value = obs.perception_data[f"{camera_name}_rgb"]
-    return np.asarray(value, dtype=np.uint8)
+    upstream_name = canonical_to_upstream_camera(camera_name)
+    for key in (f"{upstream_name}_rgb", f"{camera_name}_rgb"):
+        if hasattr(obs, key):
+            return np.asarray(getattr(obs, key), dtype=np.uint8)
+    perception_data = getattr(obs, "perception_data", None)
+    if isinstance(perception_data, dict):
+        for key in (f"{upstream_name}_rgb", f"{camera_name}_rgb"):
+            if key in perception_data:
+                return np.asarray(perception_data[key], dtype=np.uint8)
+    raise KeyError(f"Observation does not expose RGB for camera '{camera_name}'")
 
 
 def _camera_point_cloud(obs: Any, camera_name: str) -> np.ndarray:
-    value = obs.perception_data[f"{camera_name}_point_cloud"]
-    return np.asarray(value, dtype=np.float32)
+    upstream_name = canonical_to_upstream_camera(camera_name)
+    for key in (
+        f"{upstream_name}_point_cloud",
+        f"{upstream_name}_pointcloud",
+        f"{camera_name}_point_cloud",
+        f"{camera_name}_pointcloud",
+    ):
+        if hasattr(obs, key):
+            return np.asarray(getattr(obs, key), dtype=np.float32)
+    perception_data = getattr(obs, "perception_data", None)
+    if isinstance(perception_data, dict):
+        for key in (
+            f"{upstream_name}_point_cloud",
+            f"{upstream_name}_pointcloud",
+            f"{camera_name}_point_cloud",
+            f"{camera_name}_pointcloud",
+        ):
+            if key in perception_data:
+                return np.asarray(perception_data[key], dtype=np.float32)
+    raise KeyError(f"Observation does not expose point clouds for camera '{camera_name}'")
+
+
+def _camera_misc(obs: Any, camera_name: str, field_name: str) -> np.ndarray:
+    upstream_name = canonical_to_upstream_camera(camera_name)
+    misc = getattr(obs, "misc", {})
+    for key in (
+        f"{upstream_name}_camera_{field_name}",
+        f"{camera_name}_camera_{field_name}",
+        f"{upstream_name}_{field_name}",
+        f"{camera_name}_{field_name}",
+    ):
+        if key in misc:
+            return np.asarray(misc[key], dtype=np.float32)
+    raise KeyError(f"Observation misc does not expose {field_name} for camera '{camera_name}'")
 
 
 def _bimanual_proprio(obs: Any, timestep: int | None = None, episode_length: int | None = None) -> np.ndarray:
-    right = np.asarray(obs.get_low_dim_data(obs.right), dtype=np.float32)
-    left = np.asarray(obs.get_low_dim_data(obs.left), dtype=np.float32)
+    if hasattr(obs, "get_low_dim_data"):
+        try:
+            right = np.asarray(obs.get_low_dim_data("right"), dtype=np.float32)
+            left = np.asarray(obs.get_low_dim_data("left"), dtype=np.float32)
+        except Exception:
+            right = np.asarray(obs.get_low_dim_data(getattr(obs, "right")), dtype=np.float32)
+            left = np.asarray(obs.get_low_dim_data(getattr(obs, "left")), dtype=np.float32)
+    else:
+        right = np.asarray(getattr(obs.right, "get_low_dim_data")(), dtype=np.float32)
+        left = np.asarray(getattr(obs.left, "get_low_dim_data")(), dtype=np.float32)
     proprio = np.concatenate([right, left], axis=0)
     if timestep is not None and episode_length and episode_length > 1:
         time_feature = np.array(
@@ -64,11 +112,11 @@ def extract_canonical_bimanual_obs(
     camera_spec = camera_spec or RLBenchThreeCameraSpec()
     rgb = {camera: _camera_rgb(obs, camera) for camera in camera_spec.cameras}
     intrinsics = {
-        camera: np.asarray(obs.misc[f"{camera}_camera_intrinsics"], dtype=np.float32)
+        camera: _camera_misc(obs, camera, "intrinsics")
         for camera in camera_spec.cameras
     }
     extrinsics = {
-        camera: np.asarray(obs.misc[f"{camera}_camera_extrinsics"], dtype=np.float32)
+        camera: _camera_misc(obs, camera, "extrinsics")
         for camera in camera_spec.cameras
     }
     point_cloud = None

code/reveal_vla_bimanual/sim_rlbench/peract2_runner.py

@@ -29,22 +29,38 @@ def _default_nvidia_shim_root() -> Path | None:
     return candidate if candidate.exists() else None
 
 
+def resolve_upstream_root(upstream_root: Path) -> Path:
+    if (upstream_root / "train.py").exists():
+        return upstream_root
+    nested_root = upstream_root / "peract"
+    if (nested_root / "train.py").exists():
+        return nested_root
+    return upstream_root
+
+
 @dataclass
 class BenchmarkRunSpec:
     upstream_root: Path = Path("/workspace/third_party/peract_bimanual")
     demo_path: Path = Path("/workspace/data/rlbench2")
     replay_path: Path = Path("/workspace/replays/rlbench2")
     logdir: Path = Path("/workspace/logs/rlbench2")
-    method: str = "BIMANUAL_PERACT"
+    method: str = "PERACT_BC"
     tasks: tuple[str, ...] = field(default_factory=lambda: PERACT2_BIMANUAL_TASKS)
     demos: int = 100
     training_iterations: int = 40000
     seed: int = 0
     gpu: int = 0
     display: str = ":99"
+    use_xvfb: bool = True
+    xvfb_screen: str = "1280x1024x24"
     coppeliasim_root: Path = Path("/workspace/assets/coppeliasim_v4_1_0")
     camera_spec: RLBenchThreeCameraSpec = field(default_factory=default_three_camera_spec)
 
+    def _wrap_display(self, command: list[str]) -> list[str]:
+        if not self.use_xvfb:
+            return command
+        return ["xvfb-run", "-a", "-s", f"-screen 0 {self.xvfb_screen}", *command]
+
     def common_overrides(self) -> list[str]:
         task_name = "multi_3cam" if len(self.tasks) > 1 else self.tasks[0]
         overrides = [
@@ -66,7 +82,7 @@ class BenchmarkRunSpec:
 
     def train_command(self, python_executable: str | None = None) -> list[str]:
         python_executable = python_executable or sys.executable
-        return [python_executable, "train.py", *self.common_overrides()]
+        return self._wrap_display([python_executable, "train.py", *self.common_overrides()])
 
     def eval_command(
         self,
@@ -76,7 +92,7 @@ class BenchmarkRunSpec:
         python_executable: str | None = None,
     ) -> list[str]:
         python_executable = python_executable or sys.executable
-        return [
+        return self._wrap_display([
             python_executable,
             "eval.py",
             f"method={self.method}",
@@ -85,7 +101,7 @@ class BenchmarkRunSpec:
             f"eval_episodes={episodes}",
             f"cinematic_recorder.enabled={str(save_videos)}",
             *self.camera_spec.hydra_overrides(),
-        ]
+        ])
 
     def env(self) -> dict[str, str]:
         env = os.environ.copy()
@@ -117,7 +133,7 @@ class BenchmarkRunSpec:
     def run_train(self) -> subprocess.CompletedProcess[bytes]:
         return subprocess.run(
             self.train_command(),
-            cwd=self.upstream_root,
+            cwd=resolve_upstream_root(self.upstream_root),
             env=self.env(),
             check=True,
         )

code/reveal_vla_bimanual/sim_rlbench/smoke_test.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 import argparse
 import json
+import shlex
 from pathlib import Path
 
 from sim_rlbench.camera_spec import default_three_camera_spec
-from sim_rlbench.peract2_runner import BenchmarkRunSpec
+from sim_rlbench.peract2_runner import BenchmarkRunSpec, resolve_upstream_root
 
 
 def main() -> None:
@@ -20,6 +21,7 @@ def main() -> None:
         "camera_names": list(spec.cameras),
         "resolution": list(spec.resolution),
         "global_camera": spec.global_camera,
+        "resolved_upstream_root": str(resolve_upstream_root(Path(args.upstream_root))),
     }
 
     import_status = {}
@@ -39,7 +41,7 @@ def main() -> None:
             demo_path=Path(args.demo_path),
             camera_spec=spec,
         )
-        print(" ".join(run_spec.train_command()))
+        print(shlex.join(run_spec.train_command()))
 
 
 if __name__ == "__main__":