Add files using upload-large-folder tool

code/reveal_vla_bimanual/sim_rlbench/launch_smoke.py

@@ -72,7 +72,8 @@ def main() -> None:
         env.launch()
         task = env.get_task(task_cls)
         descriptions, obs = task.reset()
-        next_obs, reward, done = task.step(_noop_bimanual_action(obs))
+        noop_action = _noop_bimanual_action(obs)
+        next_obs, reward, done = task.step(noop_action)
         payload = {
             "display": args.display or os.environ.get("DISPLAY"),
             "headless": headless,
@@ -84,6 +85,8 @@ def main() -> None:
             "right_pose_shape": list(obs.right.gripper_pose.shape),
             "left_pose_shape": list(obs.left.gripper_pose.shape),
             "stepped_mode": "bimanual_noop",
+            "action_finite": bool(np.isfinite(noop_action).all()),
+            "action_dim": int(noop_action.shape[0]),
             "reward": float(reward),
             "done": bool(done),
             "front_rgb_shape_after_step": _camera_shape(next_obs, "front"),

code/reveal_vla_bimanual/sim_rlbench/task_resolver.py

@@ -8,4 +8,8 @@ def is_bimanual_task_name(task_name: str) -> bool:
 
 
 def resolve_task_class(task_name: str):
+    if task_name == "open_drawer":
+        # The dual-panda RLBench fork provides dedicated bimanual drawer tasks.
+        # Route the generic smoke/rollout task name to the compatible variant.
+        return task_file_to_task_class("right_open_drawer", bimanual=True)
     return task_file_to_task_class(task_name, bimanual=is_bimanual_task_name(task_name))

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_ablation_nodepth.yaml

@@ -0,0 +1,16 @@
+experiment_name: proxy_interaction_r3d_ablation_nodepth
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 17
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 224, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_clip224_v6_rgbd_stage3.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_clip224_v6_rgbd_stage3.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 17}
+optim: {epochs: 4, batch_size: 2, num_workers: 0, lr: 0.0003, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: true, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: true, plan_during_eval: true, support_mode_conditioning: true, planner_mode: trainable, use_depth: false, use_world_model: true, use_role_tokens: true, compute_equivariance_probe: true}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 512, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: false}
+  fusion: {hidden_dim: 512, num_cameras: 3, num_layers: 4, num_heads: 8, ff_dim: 2048, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 512, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 8, max_history_steps: 8}
+  decoder: {hidden_dim: 512, num_heads: 8, num_layers: 4, ff_dim: 2048, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 512, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 8, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 512, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 8, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 512, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 8, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0, phase: 0.1, arm_role: 0.15, support_mode: 0.1, corridor: 0.15, persistence: 0.05, disturbance: 0.05, world_model: 0.25, belief: 0.05, visibility: 0.05, clearance: 0.05, support_stability: 0.05, reocclusion: 0.05, occluder_contact: 0.05, grasp_affordance: 0.05, planner_success: 0.25, planner_risk: 0.1, planner_ranking: 0.2, proposal_reconstruction: 0.1, proposal_success: 0.15, proposal_ranking: 0.2, proposal_diversity: 0.05, role_swap_consistency: 0.05}

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_ablation_noplanner.yaml

@@ -0,0 +1,16 @@
+experiment_name: proxy_interaction_r3d_ablation_noplanner
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 17
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 96, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_v6_rgbd_stage2_dummy.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_v6_rgbd_stage2_dummy.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 17}
+optim: {epochs: 10, batch_size: 16, num_workers: 0, lr: 0.001, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: false, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: false, plan_during_eval: false, support_mode_conditioning: true, planner_mode: off, use_depth: false, use_world_model: true, use_role_tokens: true, compute_equivariance_probe: true}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 192, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: true}
+  fusion: {hidden_dim: 192, num_cameras: 3, num_layers: 2, num_heads: 4, ff_dim: 384, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 192, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 4, max_history_steps: 8}
+  decoder: {hidden_dim: 192, num_heads: 4, num_layers: 2, ff_dim: 384, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 192, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 4, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 192, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 4, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 192, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 4, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0}

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_ablation_norolesym.yaml

@@ -0,0 +1,16 @@
+experiment_name: proxy_interaction_r3d_ablation_norolesym
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 17
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 96, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_v6_rgbd_stage2_dummy.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_v6_rgbd_stage2_dummy.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 17}
+optim: {epochs: 10, batch_size: 16, num_workers: 0, lr: 0.001, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: false, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: true, plan_during_eval: true, support_mode_conditioning: true, planner_mode: trainable, use_depth: false, use_world_model: true, use_role_tokens: false, compute_equivariance_probe: false}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 192, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: true}
+  fusion: {hidden_dim: 192, num_cameras: 3, num_layers: 2, num_heads: 4, ff_dim: 384, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 192, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 4, max_history_steps: 8}
+  decoder: {hidden_dim: 192, num_heads: 4, num_layers: 2, ff_dim: 384, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 192, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 4, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 192, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 4, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 192, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 4, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0}

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_ablation_nowm.yaml

@@ -0,0 +1,16 @@
+experiment_name: proxy_interaction_r3d_ablation_nowm
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 17
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 96, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_v6_rgbd_stage2_dummy.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_v6_rgbd_stage2_dummy.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 17}
+optim: {epochs: 10, batch_size: 16, num_workers: 0, lr: 0.001, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: false, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: true, plan_during_eval: true, support_mode_conditioning: true, planner_mode: trainable, use_depth: false, use_world_model: false, use_role_tokens: true, compute_equivariance_probe: true}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 192, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: true}
+  fusion: {hidden_dim: 192, num_cameras: 3, num_layers: 2, num_heads: 4, ff_dim: 384, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 192, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 4, max_history_steps: 8}
+  decoder: {hidden_dim: 192, num_heads: 4, num_layers: 2, ff_dim: 384, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 192, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 4, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 192, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 4, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 192, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 4, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0}

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_stage1_clip.yaml

@@ -0,0 +1,72 @@
+experiment_name: proxy_interaction_r3d_stage1_clip
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 7
+init_checkpoint: /workspace/VLAarchtests/artifacts/outputs/reveal_runs/proxy_backbone_only_clip/checkpoint_best.pt
+init_strict: false
+
+data:
+  proxies: [foliage_proxy, bag_proxy, cloth_proxy]
+  resolution: 224
+  dataset_version: reveal_proxy_v6_rgbd_elastic_state
+  train_episodes_per_proxy: 48
+  val_episodes_per_proxy: 16
+  train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_clip224_v6_rgbd_stage1.pt
+  val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_clip224_v6_rgbd_stage1.pt
+  rebuild_dataset: false
+  chunk_horizon: 8
+  rollout_horizon: 5
+  history_steps: 6
+  planner_candidates: 8
+  seed: 7
+
+optim: {epochs: 4, batch_size: 2, num_workers: 4, lr: 0.0003, weight_decay: 0.0001}
+
+trainer:
+  policy_type: elastic_reveal
+  use_bf16: true
+  grad_clip_norm: 1.0
+  freeze_backbone: true
+  gradient_checkpointing: false
+  plan_during_train: true
+  plan_during_eval: true
+  support_mode_conditioning: true
+  planner_mode: trainable
+  use_depth: false
+  use_world_model: true
+  use_role_tokens: true
+  compute_equivariance_probe: true
+
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 512, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: false}
+  fusion: {hidden_dim: 512, num_cameras: 3, num_layers: 4, num_heads: 8, ff_dim: 2048, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 512, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 8, max_history_steps: 8}
+  decoder: {hidden_dim: 512, num_heads: 8, num_layers: 4, ff_dim: 2048, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 512, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 8, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 512, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 8, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 512, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 8, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+
+loss_weights:
+  action: 1.0
+  phase: 0.1
+  arm_role: 0.15
+  support_mode: 0.1
+  corridor: 0.15
+  persistence: 0.05
+  disturbance: 0.05
+  world_model: 0.2
+  belief: 0.05
+  visibility: 0.05
+  clearance: 0.05
+  support_stability: 0.05
+  reocclusion: 0.05
+  occluder_contact: 0.05
+  grasp_affordance: 0.05
+  planner_success: 0.25
+  planner_risk: 0.1
+  planner_ranking: 0.2
+  proposal_reconstruction: 0.1
+  proposal_success: 0.15
+  proposal_ranking: 0.2
+  proposal_diversity: 0.05
+  role_swap_consistency: 0.05

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_stage1_dummy.yaml

@@ -0,0 +1,75 @@
+experiment_name: proxy_interaction_r3d_stage1_dummy
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 13
+
+data:
+  proxies: [foliage_proxy, bag_proxy, cloth_proxy]
+  resolution: 96
+  dataset_version: reveal_proxy_v6_rgbd_elastic_state
+  train_episodes_per_proxy: 48
+  val_episodes_per_proxy: 16
+  train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_v6_rgbd_stage1_dummy.pt
+  val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_v6_rgbd_stage1_dummy.pt
+  rebuild_dataset: false
+  chunk_horizon: 8
+  rollout_horizon: 5
+  history_steps: 6
+  planner_candidates: 8
+  seed: 13
+
+optim:
+  epochs: 4
+  batch_size: 16
+  num_workers: 4
+  lr: 0.001
+  weight_decay: 0.0001
+
+trainer:
+  policy_type: elastic_reveal
+  use_bf16: false
+  grad_clip_norm: 1.0
+  freeze_backbone: true
+  gradient_checkpointing: false
+  plan_during_train: true
+  plan_during_eval: true
+  support_mode_conditioning: true
+  planner_mode: trainable
+  use_depth: false
+  use_world_model: true
+  use_role_tokens: true
+  compute_equivariance_probe: true
+
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 192, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: true}
+  fusion: {hidden_dim: 192, num_cameras: 3, num_layers: 2, num_heads: 4, ff_dim: 384, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 192, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 4, max_history_steps: 8}
+  decoder: {hidden_dim: 192, num_heads: 4, num_layers: 2, ff_dim: 384, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 192, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 4, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 192, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 4, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 192, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 4, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+
+loss_weights:
+  action: 1.0
+  phase: 0.15
+  arm_role: 0.2
+  support_mode: 0.15
+  corridor: 0.2
+  persistence: 0.1
+  disturbance: 0.1
+  world_model: 0.25
+  belief: 0.05
+  visibility: 0.05
+  clearance: 0.05
+  support_stability: 0.05
+  reocclusion: 0.05
+  occluder_contact: 0.05
+  grasp_affordance: 0.05
+  planner_success: 0.2
+  planner_risk: 0.1
+  planner_ranking: 0.1
+  proposal_reconstruction: 0.2
+  proposal_success: 0.1
+  proposal_ranking: 0.1
+  proposal_diversity: 0.05
+  role_swap_consistency: 0.05

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_stage2_clip.yaml

@@ -0,0 +1,18 @@
+experiment_name: proxy_interaction_r3d_stage2_clip
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 11
+init_checkpoint: /workspace/VLAarchtests/artifacts/outputs/reveal_runs/proxy_backbone_only_clip/checkpoint_best.pt
+init_strict: false
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 224, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_clip224_v6_rgbd_stage2.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_clip224_v6_rgbd_stage2.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 11}
+optim: {epochs: 4, batch_size: 2, num_workers: 4, lr: 0.0003, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: true, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: true, plan_during_eval: true, support_mode_conditioning: true, planner_mode: trainable, use_depth: false, use_world_model: true, use_role_tokens: true, compute_equivariance_probe: true}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 512, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: false}
+  fusion: {hidden_dim: 512, num_cameras: 3, num_layers: 4, num_heads: 8, ff_dim: 2048, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 512, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 8, max_history_steps: 8}
+  decoder: {hidden_dim: 512, num_heads: 8, num_layers: 4, ff_dim: 2048, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 512, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 8, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 512, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 8, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 512, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 8, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0, phase: 0.1, arm_role: 0.15, support_mode: 0.1, corridor: 0.15, persistence: 0.05, disturbance: 0.05, world_model: 0.25, belief: 0.05, visibility: 0.05, clearance: 0.05, support_stability: 0.05, reocclusion: 0.05, occluder_contact: 0.05, grasp_affordance: 0.05, planner_success: 0.25, planner_risk: 0.1, planner_ranking: 0.2, proposal_reconstruction: 0.1, proposal_success: 0.15, proposal_ranking: 0.2, proposal_diversity: 0.05, role_swap_consistency: 0.05}

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_stage2_dummy.yaml

@@ -0,0 +1,17 @@
+experiment_name: proxy_interaction_r3d_stage2_dummy
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 21
+defaults: []
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 96, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_v6_rgbd_stage2_dummy.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_v6_rgbd_stage2_dummy.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 21}
+optim: {epochs: 4, batch_size: 16, num_workers: 4, lr: 0.001, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: false, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: true, plan_during_eval: true, support_mode_conditioning: true, planner_mode: trainable, use_depth: false, use_world_model: true, use_role_tokens: true, compute_equivariance_probe: true}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 192, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: true}
+  fusion: {hidden_dim: 192, num_cameras: 3, num_layers: 2, num_heads: 4, ff_dim: 384, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 192, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 4, max_history_steps: 8}
+  decoder: {hidden_dim: 192, num_heads: 4, num_layers: 2, ff_dim: 384, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 192, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 4, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 192, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 4, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 192, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 4, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0, phase: 0.15, arm_role: 0.2, support_mode: 0.15, corridor: 0.2, persistence: 0.1, disturbance: 0.1, world_model: 0.3, belief: 0.05, visibility: 0.05, clearance: 0.05, support_stability: 0.05, reocclusion: 0.05, occluder_contact: 0.05, grasp_affordance: 0.05, planner_success: 0.2, planner_risk: 0.1, planner_ranking: 0.1, proposal_reconstruction: 0.2, proposal_success: 0.1, proposal_ranking: 0.1, proposal_diversity: 0.05, role_swap_consistency: 0.05}

code/reveal_vla_bimanual/train/configs/proxy_interaction_r3d_stage3_clip_rgbd.yaml

@@ -0,0 +1,18 @@
+experiment_name: proxy_interaction_r3d_stage3_clip_rgbd
+output_dir: /workspace/VLAarchtests/artifacts/outputs/r3d
+device: cuda
+seed: 17
+init_checkpoint: /workspace/VLAarchtests/artifacts/outputs/reveal_runs/proxy_backbone_only_clip/checkpoint_best.pt
+init_strict: false
+data: {proxies: [foliage_proxy, bag_proxy, cloth_proxy], resolution: 224, dataset_version: reveal_proxy_v6_rgbd_elastic_state, train_episodes_per_proxy: 48, val_episodes_per_proxy: 16, train_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_train_clip224_v6_rgbd_stage3.pt, val_dataset_path: /workspace/VLAarchtests/artifacts/data/reveal_proxy/proxy_val_clip224_v6_rgbd_stage3.pt, rebuild_dataset: false, chunk_horizon: 8, rollout_horizon: 5, history_steps: 6, planner_candidates: 8, seed: 17}
+optim: {epochs: 4, batch_size: 2, num_workers: 4, lr: 0.0003, weight_decay: 0.0001}
+trainer: {policy_type: elastic_reveal, use_bf16: true, grad_clip_norm: 1.0, freeze_backbone: true, gradient_checkpointing: false, plan_during_train: true, plan_during_eval: true, support_mode_conditioning: true, planner_mode: trainable, use_depth: true, use_world_model: true, use_role_tokens: true, compute_equivariance_probe: true}
+policy:
+  backbone: {model_name: openai/clip-vit-base-patch32, hidden_dim: 512, max_text_tokens: 32, freeze_backbone: true, gradient_checkpointing: false, use_dummy_backbone: false}
+  fusion: {hidden_dim: 512, num_cameras: 3, num_layers: 4, num_heads: 8, ff_dim: 2048, dropout: 0.1, proprio_dim: 32, proprio_tokens: 1}
+  memory: {hidden_dim: 512, action_dim: 14, history_steps: 6, scene_history_steps: 3, belief_history_steps: 8, num_layers: 2, dropout: 0.1, memory_bank_size: 4, scene_bank_size: 2, belief_bank_size: 2, num_heads: 8, max_history_steps: 8}
+  decoder: {hidden_dim: 512, num_heads: 8, num_layers: 4, ff_dim: 2048, dropout: 0.1, chunk_size: 8, action_dim: 14, arm_action_dim: 7, num_candidates: 8, num_phases: 5, num_arm_roles: 4, num_proposal_modes: 6, planner_top_k: 4}
+  reveal_head: {hidden_dim: 512, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, belief_map_size: 32, field_size: 16, num_heads: 8, predict_belief_map: true, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8}
+  world_model: {hidden_dim: 512, action_dim: 14, num_support_modes: 3, num_approach_templates: 32, rollout_horizon: 5, field_size: 16, num_heads: 8, num_phases: 5, num_arm_roles: 4, num_interaction_tokens: 8, belief_map_size: 32, predict_belief_map: true, scene_bank_size: 2, belief_bank_size: 2}
+  planner: {hidden_dim: 512, num_candidates: 8, action_dim: 14, num_support_modes: 3, utility_margin: 0.1, num_heads: 8, num_layers: 2, num_phases: 5, num_arm_roles: 4, top_k: 4}
+loss_weights: {action: 1.0, phase: 0.1, arm_role: 0.15, support_mode: 0.1, corridor: 0.15, persistence: 0.05, disturbance: 0.05, world_model: 0.25, belief: 0.05, visibility: 0.05, clearance: 0.05, support_stability: 0.05, reocclusion: 0.05, occluder_contact: 0.05, grasp_affordance: 0.05, planner_success: 0.25, planner_risk: 0.1, planner_ranking: 0.2, proposal_reconstruction: 0.1, proposal_success: 0.15, proposal_ranking: 0.2, proposal_diversity: 0.05, role_swap_consistency: 0.05}

code/reveal_vla_bimanual/train/losses.py

@@ -18,12 +18,20 @@ class LossWeights:
     disturbance: float = 0.05
     world_model: float = 0.1
     belief: float = 0.05
+    visibility: float = 0.05
+    clearance: float = 0.05
+    support_stability: float = 0.05
+    reocclusion: float = 0.05
+    occluder_contact: float = 0.05
+    grasp_affordance: float = 0.05
     planner_success: float = 0.1
     planner_risk: float = 0.05
     planner_ranking: float = 0.05
     proposal_reconstruction: float = 0.1
     proposal_success: float = 0.05
     proposal_ranking: float = 0.05
+    proposal_diversity: float = 0.05
+    role_swap_consistency: float = 0.05
 
 
 def chunk_bc_loss(pred_actions: Tensor, target_actions: Tensor, mask: Tensor | None = None) -> Tensor:
@@ -62,6 +70,46 @@ def _role_targets_like(arm_role_logits: Tensor) -> Tensor:
     return role_target.view(*expand_shape).expand(*arm_role_logits.shape[:-1])
 
 
+def swap_arm_actions(action_chunk: Tensor) -> Tensor:
+    midpoint = action_chunk.shape[-1] // 2
+    return torch.cat([action_chunk[..., midpoint:], action_chunk[..., :midpoint]], dim=-1)
+
+
+def permutation_invariant_role_loss(arm_role_logits: Tensor) -> Tensor:
+    role_target = _role_targets_like(arm_role_logits)
+    swapped_target = role_target.flip(-1)
+    flat_logits = arm_role_logits.reshape(-1, arm_role_logits.shape[-1])
+    loss_a = F.cross_entropy(flat_logits, role_target.reshape(-1), reduction="none").view(*role_target.shape)
+    loss_b = F.cross_entropy(flat_logits, swapped_target.reshape(-1), reduction="none").view(*role_target.shape)
+    return torch.minimum(loss_a.sum(dim=-1), loss_b.sum(dim=-1)).mean()
+
+
+def role_swap_consistency_loss(pred: Tensor, target: Tensor) -> Tensor:
+    return F.smooth_l1_loss(pred, target)
+
+
+def proposal_diversity_loss(proposal_candidates: Tensor, minimum_distance: float = 0.05) -> Tensor:
+    if proposal_candidates.ndim != 4 or proposal_candidates.shape[1] <= 1:
+        return proposal_candidates.new_tensor(0.0)
+    flat = proposal_candidates.flatten(start_dim=2)
+    distances = torch.cdist(flat, flat, p=1)
+    eye = torch.eye(distances.shape[-1], device=distances.device, dtype=torch.bool).unsqueeze(0)
+    valid = (~eye).expand(distances.shape[0], -1, -1)
+    if not valid.any():
+        return proposal_candidates.new_tensor(0.0)
+    return torch.relu(minimum_distance - distances[valid]).mean()
+
+
+def _resize_like(target: Tensor, prediction: Tensor) -> Tensor:
+    if target.shape == prediction.shape:
+        return target
+    if target.ndim == prediction.ndim == 4:
+        return F.interpolate(target.float(), size=prediction.shape[-2:], mode="bilinear", align_corners=False)
+    if target.ndim == 3 and prediction.ndim == 4:
+        return F.interpolate(target.unsqueeze(1).float(), size=prediction.shape[-2:], mode="bilinear", align_corners=False)
+    return target
+
+
 def reveal_state_loss(pred: dict[str, Tensor], target: dict[str, Tensor], weights: LossWeights) -> dict[str, Tensor]:
     losses = {}
     if "phase_logits" in pred:
@@ -75,11 +123,7 @@ def reveal_state_loss(pred: dict[str, Tensor], target: dict[str, Tensor], weight
     else:
         losses["phase"] = pred["support_mode_logits"].new_tensor(0.0)
     if "arm_role_logits" in pred:
-        role_target = _role_targets_like(pred["arm_role_logits"])
-        role_ce = F.cross_entropy(
-            pred["arm_role_logits"].reshape(-1, pred["arm_role_logits"].shape[-1]),
-            role_target.reshape(-1),
-        )
+        role_ce = permutation_invariant_role_loss(pred["arm_role_logits"])
         role_probs = pred["arm_role_logits"].softmax(dim=-1)
         role_gap = torch.mean(torch.abs(role_probs[:, 0] - role_probs[:, 1]), dim=-1)
         role_separation = torch.relu(0.25 - role_gap).mean()
@@ -95,11 +139,50 @@ def reveal_state_loss(pred: dict[str, Tensor], target: dict[str, Tensor], weight
     losses["persistence"] = F.mse_loss(pred["persistence_horizon"], target["persistence_horizon"].float())
     losses["disturbance"] = F.mse_loss(pred["disturbance_cost"], target["disturbance_cost"].float())
     if "belief_map" in pred and "belief_map" in target:
-        losses["belief"] = F.binary_cross_entropy_with_logits(pred["belief_map"], target["belief_map"].float())
+        losses["belief"] = F.binary_cross_entropy_with_logits(pred["belief_map"], _resize_like(target["belief_map"].float(), pred["belief_map"]))
     else:
         losses["belief"] = pred["support_mode_logits"].new_tensor(0.0)
+    if "visibility_field" in pred and "visibility_map" in target:
+        losses["visibility"] = F.binary_cross_entropy_with_logits(
+            pred["visibility_field"],
+            _resize_like(target["visibility_map"].float(), pred["visibility_field"]),
+        )
+    else:
+        losses["visibility"] = pred["support_mode_logits"].new_tensor(0.0)
+    if "clearance_field" in pred and "clearance_map" in target:
+        losses["clearance"] = F.binary_cross_entropy_with_logits(
+            pred["clearance_field"],
+            _resize_like(target["clearance_map"].float(), pred["clearance_field"]),
+        )
+    else:
+        losses["clearance"] = pred["support_mode_logits"].new_tensor(0.0)
+    if "support_stability_field" in pred and "support_stability_map" in target:
+        losses["support_stability"] = F.binary_cross_entropy_with_logits(
+            pred["support_stability_field"],
+            _resize_like(target["support_stability_map"].float(), pred["support_stability_field"]),
+        )
+    else:
+        losses["support_stability"] = pred["support_mode_logits"].new_tensor(0.0)
+    if "occluder_contact_field" in pred and "occluder_contact_map" in target:
+        losses["occluder_contact"] = F.binary_cross_entropy_with_logits(
+            pred["occluder_contact_field"],
+            _resize_like(target["occluder_contact_map"].float(), pred["occluder_contact_field"]),
+        )
+    else:
+        losses["occluder_contact"] = pred["support_mode_logits"].new_tensor(0.0)
+    if "grasp_affordance_field" in pred and "grasp_affordance_map" in target:
+        losses["grasp_affordance"] = F.binary_cross_entropy_with_logits(
+            pred["grasp_affordance_field"],
+            _resize_like(target["grasp_affordance_map"].float(), pred["grasp_affordance_field"]),
+        )
+    else:
+        losses["grasp_affordance"] = pred["support_mode_logits"].new_tensor(0.0)
     if "reocclusion_logit" in pred and "corridor_feasible" in target:
-        target_reocclusion = 1.0 - target["corridor_feasible"].float().amax(dim=-1)
+        target_reocclusion = target.get("reocclusion_target")
+        if target_reocclusion is None:
+            target_reocclusion = 1.0 - target["corridor_feasible"].float().amax(dim=-1)
+        if target_reocclusion.ndim < pred["reocclusion_logit"].ndim:
+            target_reocclusion = target_reocclusion.unsqueeze(-1).expand_as(pred["reocclusion_logit"])
         losses["reocclusion"] = F.binary_cross_entropy_with_logits(pred["reocclusion_logit"], target_reocclusion)
     else:
         losses["reocclusion"] = pred["support_mode_logits"].new_tensor(0.0)
@@ -157,11 +240,30 @@ def world_model_rollout_consistency_loss(pred_rollout: dict[str, Tensor], target
             phase_target.reshape(-1),
         )
     if "arm_role_logits" in pred_rollout:
-        role_target = _role_targets_like(pred_rollout["arm_role_logits"])
-        loss = loss + 0.25 * F.cross_entropy(
-            pred_rollout["arm_role_logits"].reshape(-1, pred_rollout["arm_role_logits"].shape[-1]),
-            role_target.reshape(-1),
-        )
+        loss = loss + 0.25 * permutation_invariant_role_loss(pred_rollout["arm_role_logits"])
+    optional_pairs = (
+        ("target_belief_field", "belief_map", "rollout_belief_map"),
+        ("visibility_field", "visibility_map", "rollout_visibility_map"),
+        ("clearance_field", "clearance_map", "rollout_clearance_map"),
+        ("support_stability_field", "support_stability_map", "rollout_support_stability"),
+        ("occluder_contact_field", "occluder_contact_map", "rollout_occluder_contact_map"),
+        ("grasp_affordance_field", "grasp_affordance_map", "rollout_grasp_affordance_map"),
+        ("reocclusion_field", "reocclusion_map", "rollout_reocclusion_target"),
+    )
+    for pred_key, _, target_key in optional_pairs:
+        if pred_key not in pred_rollout or target_key not in target_rollout:
+            continue
+        target_value = _expand_target(target_rollout[target_key][..., :horizon, ...])
+        pred_value = pred_rollout[pred_key][..., :horizon, :, :, :] if pred_rollout[pred_key].ndim >= 6 else pred_rollout[pred_key][..., :horizon, :, :]
+        if target_value.ndim + 1 == pred_value.ndim:
+            target_value = target_value.unsqueeze(-3)
+        if pred_key in {"support_stability_field", "reocclusion_field"} and target_value.ndim == pred_value.ndim - 1:
+            target_value = target_value.unsqueeze(-1).unsqueeze(-1)
+        if pred_key == "support_stability_field" and target_value.ndim == pred_value.ndim - 2:
+            target_value = target_value.unsqueeze(-1).unsqueeze(-1)
+        if pred_key == "reocclusion_field" and target_value.ndim == pred_value.ndim - 2:
+            target_value = target_value.unsqueeze(-1).unsqueeze(-1)
+        loss = loss + 0.1 * F.binary_cross_entropy_with_logits(pred_value, target_value.float())
     return loss
 
 
@@ -192,7 +294,12 @@ def compute_total_loss(
             + weights.persistence * reveal_losses["persistence"]
             + weights.disturbance * reveal_losses["disturbance"]
             + weights.belief * reveal_losses["belief"]
-            + 0.02 * reveal_losses["reocclusion"]
+            + weights.visibility * reveal_losses["visibility"]
+            + weights.clearance * reveal_losses["clearance"]
+            + weights.support_stability * reveal_losses["support_stability"]
+            + weights.occluder_contact * reveal_losses["occluder_contact"]
+            + weights.grasp_affordance * reveal_losses["grasp_affordance"]
+            + weights.reocclusion * reveal_losses["reocclusion"]
             + 0.01 * reveal_losses["uncertainty"]
         )
 
@@ -207,6 +314,28 @@ def compute_total_loss(
                 "disturbance_cost": batch["candidate_rollout_disturbance_cost"],
                 "action_chunk": batch["candidate_action_chunks"],
             }
+            for optional_key in (
+                "candidate_rollout_belief_map",
+                "candidate_rollout_visibility_map",
+                "candidate_rollout_clearance_map",
+                "candidate_rollout_support_stability",
+                "candidate_rollout_reocclusion_target",
+                "candidate_rollout_occluder_contact_map",
+                "candidate_rollout_grasp_affordance_map",
+            ):
+                if optional_key in batch:
+                    rollout_target[optional_key.replace("candidate_", "")] = batch[optional_key]
+            planner_indices = model_output.get("planner_topk_indices")
+            if planner_indices is not None:
+                for key, value in list(rollout_target.items()):
+                    if isinstance(value, Tensor) and value.ndim >= 2 and value.shape[1] >= planner_indices.shape[1]:
+                        expand_indices = planner_indices
+                        while expand_indices.ndim < value.ndim:
+                            expand_indices = expand_indices.unsqueeze(-1)
+                        rollout_target[key] = value.gather(
+                            1,
+                            expand_indices.expand(-1, -1, *value.shape[2:]),
+                        )
         else:
             rollout_target = {
                 "support_mode": batch["rollout_support_mode"],
@@ -215,6 +344,17 @@ def compute_total_loss(
                 "disturbance_cost": batch["rollout_disturbance_cost"],
                 "action_chunk": batch["action_chunk"],
             }
+            for optional_key in (
+                "rollout_belief_map",
+                "rollout_visibility_map",
+                "rollout_clearance_map",
+                "rollout_support_stability",
+                "rollout_reocclusion_target",
+                "rollout_occluder_contact_map",
+                "rollout_grasp_affordance_map",
+            ):
+                if optional_key in batch:
+                    rollout_target[optional_key] = batch[optional_key]
         world_model_loss = world_model_rollout_consistency_loss(
             model_output["planned_rollout"],
             rollout_target,
@@ -234,8 +374,16 @@ def compute_total_loss(
                 1.0,
             )
         utility_target = batch.get("candidate_utility")
+        planner_indices = model_output.get("planner_topk_indices")
+        if planner_indices is not None and success_target.shape[1] != model_output["planner_success_logits"].shape[1]:
+            success_target = success_target.gather(1, planner_indices)
+            risk_target = risk_target.gather(1, planner_indices)
+            if utility_target is not None:
+                utility_target = utility_target.gather(1, planner_indices)
         if utility_target is None:
             utility_target = success_target - risk_target
+        elif "utility_structured" in model_output:
+            utility_target = 0.5 * utility_target + 0.5 * model_output["utility_structured"].detach()
         success_loss = F.binary_cross_entropy_with_logits(model_output["planner_success_logits"], success_target)
         risk_loss = F.mse_loss(model_output["planner_risk_values"], risk_target.float())
         pred_diff = model_output["planner_scores"].unsqueeze(-1) - model_output["planner_scores"].unsqueeze(-2)
@@ -322,5 +470,22 @@ def compute_total_loss(
         losses["proposal_success"] = model_output["action_mean"].new_tensor(0.0)
         losses["proposal_ranking"] = model_output["action_mean"].new_tensor(0.0)
 
+    if "proposal_candidates" in model_output:
+        diversity_loss = proposal_diversity_loss(model_output["proposal_candidates"])
+        losses["proposal_diversity"] = diversity_loss
+        total = total + weights.proposal_diversity * diversity_loss
+    else:
+        losses["proposal_diversity"] = model_output["action_mean"].new_tensor(0.0)
+
+    if "equivariance_probe_action_mean" in model_output and "equivariance_target_action_mean" in model_output:
+        swap_loss = role_swap_consistency_loss(
+            model_output["equivariance_probe_action_mean"],
+            model_output["equivariance_target_action_mean"].detach(),
+        )
+        losses["role_swap_consistency"] = swap_loss
+        total = total + weights.role_swap_consistency * swap_loss
+    else:
+        losses["role_swap_consistency"] = model_output["action_mean"].new_tensor(0.0)
+
     losses["total"] = total
     return losses

code/reveal_vla_bimanual/train/run_experiment.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import time
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
@@ -102,6 +103,7 @@ def _build_dataset_from_config(data_cfg: Any, split: str) -> dict[str, Any]:
         rollout_horizon=int(data_cfg.rollout_horizon),
         history_steps=int(data_cfg.get("history_steps", 2)),
         planner_candidates=int(data_cfg.get("planner_candidates", 4)),
+        dataset_version=str(data_cfg.get("dataset_version", "reveal_proxy_v5_noleak_actionhist")),
     )
     if dataset_path:
         save_teacher_dataset(dataset_path, dataset_bundle)
@@ -145,6 +147,16 @@ def main() -> None:
 
     train_bundle = _build_dataset_from_config(cfg.data, "train")
     val_bundle = _build_dataset_from_config(cfg.data, "val")
+    print(
+        json.dumps(
+            {
+                "phase": "datasets_ready",
+                "train_samples": len(train_bundle["samples"]),
+                "val_samples": len(val_bundle["samples"]),
+            }
+        ),
+        flush=True,
+    )
     train_loader = _make_loader(
         train_bundle["samples"],
         resolution=int(cfg.data.resolution),
@@ -159,6 +171,7 @@ def main() -> None:
         shuffle=False,
         num_workers=int(cfg.optim.num_workers),
     )
+    print(json.dumps({"phase": "loaders_ready"}), flush=True)
 
     policy_config = _policy_config_from_omega(cfg.policy)
     trainer_config = _trainer_config_from_omega(cfg.trainer)
@@ -173,8 +186,13 @@ def main() -> None:
     history: list[dict[str, Any]] = []
     best_val = float("inf")
     best_checkpoint = output_dir / "checkpoint_best.pt"
+    train_start_time = time.time()
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+    print(json.dumps({"phase": "training_start", "epochs": int(cfg.optim.epochs)}), flush=True)
 
     for epoch in range(int(cfg.optim.epochs)):
+        epoch_start_time = time.time()
         model.train()
         train_losses: list[dict[str, float]] = []
         for batch in train_loader:
@@ -200,6 +218,22 @@ def main() -> None:
                     forward_kwargs["support_mode_conditioning"] = trainer_config.support_mode_conditioning
                     if "candidate_action_chunks" in moved:
                         forward_kwargs["candidate_chunks_override"] = moved["candidate_action_chunks"]
+                if trainer_config.policy_type == "elastic_reveal":
+                    forward_kwargs["depths"] = moved.get("depths")
+                    forward_kwargs["depth_valid"] = moved.get("depth_valid")
+                    forward_kwargs["camera_intrinsics"] = moved.get("camera_intrinsics")
+                    forward_kwargs["camera_extrinsics"] = moved.get("camera_extrinsics")
+                    forward_kwargs["history_depths"] = moved.get("history_depths")
+                    forward_kwargs["history_depth_valid"] = moved.get("history_depth_valid")
+                    forward_kwargs["use_depth"] = trainer_config.use_depth
+                    forward_kwargs["use_world_model"] = trainer_config.use_world_model
+                    forward_kwargs["use_planner"] = planner_enabled(trainer_config, during_eval=True)
+                    forward_kwargs["use_role_tokens"] = trainer_config.use_role_tokens
+                    forward_kwargs["history_steps_override"] = trainer_config.history_steps_override
+                    forward_kwargs["compute_equivariance_probe"] = False
+                elif trainer_config.policy_type == "interaction_state":
+                    forward_kwargs["use_role_tokens"] = trainer_config.use_role_tokens
+                    forward_kwargs["history_steps_override"] = trainer_config.history_steps_override
                 output = model(**forward_kwargs)
                 losses = compute_total_loss(output, moved, weights=loss_weights)
                 val_losses.append({key: float(value.detach().cpu()) for key, value in losses.items()})
@@ -212,6 +246,18 @@ def main() -> None:
             "val": val_summary,
         }
         history.append(record)
+        print(
+            json.dumps(
+                {
+                    "phase": "epoch_complete",
+                    "epoch": epoch,
+                    "train_total": train_summary["total"],
+                    "val_total": val_summary["total"],
+                    "epoch_time_sec": float(time.time() - epoch_start_time),
+                }
+            ),
+            flush=True,
+        )
         if val_summary["total"] <= best_val:
             best_val = val_summary["total"]
             torch.save(
@@ -231,12 +277,20 @@ def main() -> None:
     metrics_path = output_dir / "metrics.json"
     metrics_path.write_text(json.dumps(history, indent=2), encoding="utf-8")
     (output_dir / "config_resolved.yaml").write_text(OmegaConf.to_yaml(cfg, resolve=True), encoding="utf-8")
+    train_time_sec = float(time.time() - train_start_time)
+    peak_gpu_memory_mb = (
+        float(torch.cuda.max_memory_allocated(device) / (1024.0 * 1024.0))
+        if device.type == "cuda"
+        else 0.0
+    )
     summary = {
         "experiment_name": cfg.experiment_name,
         "device": str(device),
         "best_checkpoint": str(best_checkpoint),
         "final_train_total": history[-1]["train"]["total"],
         "final_val_total": history[-1]["val"]["total"],
+        "train_time_sec": train_time_sec,
+        "peak_gpu_memory_mb": peak_gpu_memory_mb,
         "num_train_samples": len(train_bundle["samples"]),
         "num_val_samples": len(val_bundle["samples"]),
         "planner_mode": trainer_config.planner_mode,

code/reveal_vla_bimanual/train/smoke_checks.py

@@ -17,7 +17,7 @@ from models.planner import PlannerConfig
 from models.policy import PolicyConfig
 from models.reveal_head import RevealHeadConfig
 from models.world_model import RevealWMConfig
-from sim_reveal.dataset import dataset_from_bundle, load_teacher_dataset
+from sim_reveal.dataset import RGBD_PROXY_DATASET_VERSION, collect_teacher_dataset, dataset_from_bundle, load_teacher_dataset
 from train.losses import LossWeights, compute_total_loss
 from train.trainer import TrainerConfig, apply_planner_mode, build_policy, planner_enabled
 
@@ -68,7 +68,7 @@ def _small_policy_config() -> PolicyConfig:
             num_layers=2,
             ff_dim=128,
             dropout=0.1,
-            chunk_size=4,
+            chunk_size=8,
             action_dim=14,
             arm_action_dim=7,
             num_candidates=4,
@@ -114,9 +114,9 @@ def _small_policy_config() -> PolicyConfig:
     )
 
 
-def _trainer_config(planner_mode: str = "trainable") -> TrainerConfig:
+def _trainer_config(planner_mode: str = "trainable", policy_type: str = "interaction_state") -> TrainerConfig:
     return TrainerConfig(
-        policy_type="interaction_state",
+        policy_type=policy_type,
         use_bf16=False,
         grad_clip_norm=1.0,
         freeze_backbone=True,
@@ -226,6 +226,72 @@ def main() -> None:
     )
     proxy_optimizer.step()
 
+    rgbd_bundle = collect_teacher_dataset(
+        episodes_per_proxy=1,
+        resolution=32,
+        history_steps=2,
+        planner_candidates=4,
+        dataset_version=RGBD_PROXY_DATASET_VERSION,
+    )
+    rgbd_dataset = dataset_from_bundle(rgbd_bundle, resolution=int(rgbd_bundle["resolution"]))
+    rgbd_loader = DataLoader(rgbd_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)
+    rgbd_batch = _move_batch_to_device(next(iter(rgbd_loader)), device)
+    elastic_trainer_config = _trainer_config(planner_mode="trainable", policy_type="elastic_reveal")
+    elastic_model = build_policy(policy_config, elastic_trainer_config).to(device)
+    elastic_model.eval()
+    with torch.no_grad():
+        rgb_only_outputs = elastic_model(
+            images=rgbd_batch["images"],
+            proprio=rgbd_batch["proprio"],
+            texts=rgbd_batch["texts"],
+            history_images=rgbd_batch.get("history_images"),
+            history_proprio=rgbd_batch.get("history_proprio"),
+            history_actions=rgbd_batch.get("history_actions"),
+            plan=True,
+            use_depth=False,
+        )
+        rgbd_outputs = elastic_model(
+            images=rgbd_batch["images"],
+            depths=rgbd_batch["depths"],
+            depth_valid=rgbd_batch["depth_valid"],
+            camera_intrinsics=rgbd_batch["camera_intrinsics"],
+            camera_extrinsics=rgbd_batch["camera_extrinsics"],
+            proprio=rgbd_batch["proprio"],
+            texts=rgbd_batch["texts"],
+            history_images=rgbd_batch.get("history_images"),
+            history_depths=rgbd_batch.get("history_depths"),
+            history_depth_valid=rgbd_batch.get("history_depth_valid"),
+            history_proprio=rgbd_batch.get("history_proprio"),
+            history_actions=rgbd_batch.get("history_actions"),
+            plan=True,
+            use_depth=True,
+            compute_equivariance_probe=True,
+        )
+        noplanner_outputs = elastic_model(
+            images=rgbd_batch["images"],
+            depths=rgbd_batch["depths"],
+            depth_valid=rgbd_batch["depth_valid"],
+            proprio=rgbd_batch["proprio"],
+            texts=rgbd_batch["texts"],
+            history_images=rgbd_batch.get("history_images"),
+            history_depths=rgbd_batch.get("history_depths"),
+            history_depth_valid=rgbd_batch.get("history_depth_valid"),
+            history_proprio=rgbd_batch.get("history_proprio"),
+            history_actions=rgbd_batch.get("history_actions"),
+            plan=True,
+            use_depth=True,
+            use_planner=False,
+        )
+    dataset_v6_required = (
+        "depths",
+        "depth_valid",
+        "visibility_map",
+        "clearance_map",
+        "support_stability",
+        "reocclusion_target",
+    )
+    missing_v6_fields = [key for key in dataset_v6_required if key not in rgbd_batch]
+
     rlbench_trainer_config = _trainer_config(planner_mode="proxy_pretrained")
     rlbench_model = build_policy(policy_config, rlbench_trainer_config).to(device)
     frozen_modules = apply_planner_mode(rlbench_model, rlbench_trainer_config)
@@ -283,6 +349,15 @@ def main() -> None:
             "planner_enabled_for_eval": planner_gate,
             "frozen_modules": frozen_modules,
         },
+        "elastic_reveal": {
+            "rgb_only_candidate_shape": list(rgb_only_outputs["candidate_chunks"].shape),
+            "rgbd_candidate_shape": list(rgbd_outputs["candidate_chunks"].shape),
+            "rgbd_topk_shape": list(rgbd_outputs["planner_topk_indices"].shape),
+            "rgbd_rollout_shape": list(rgbd_outputs["planned_rollout"]["target_belief_field"].shape),
+            "noplanner_chunk_shape": list(noplanner_outputs["planned_chunk"].shape),
+            "equivariance_probe_shape": list(rgbd_outputs["equivariance_probe_action_mean"].shape),
+            "dataset_v6_missing_fields": missing_v6_fields,
+        },
         "policy_config": asdict(policy_config),
     }
 

code/reveal_vla_bimanual/train/trainer.py

@@ -6,7 +6,13 @@ from typing import Sequence
 import torch
 from torch import Tensor, nn
 
-from models.policy import BackboneOnlyPolicy, InteractionBimanualPolicy, PolicyConfig, RevealBimanualPolicy
+from models.policy import (
+    BackboneOnlyPolicy,
+    ElasticRevealBimanualPolicy,
+    InteractionBimanualPolicy,
+    PolicyConfig,
+    RevealBimanualPolicy,
+)
 from train.losses import LossWeights, compute_total_loss
 
 
@@ -21,11 +27,18 @@ class TrainerConfig:
     plan_during_eval: bool = True
     support_mode_conditioning: bool = True
     planner_mode: str = "trainable"
+    use_depth: bool = True
+    use_world_model: bool = True
+    use_role_tokens: bool = True
+    history_steps_override: int | None = None
+    compute_equivariance_probe: bool = True
 
 
 def build_policy(config: PolicyConfig, trainer_config: TrainerConfig) -> nn.Module:
     config.backbone.freeze_backbone = trainer_config.freeze_backbone
     config.backbone.gradient_checkpointing = trainer_config.gradient_checkpointing
+    if trainer_config.policy_type == "elastic_reveal":
+        return ElasticRevealBimanualPolicy(config)
     if trainer_config.policy_type == "reveal_state":
         return RevealBimanualPolicy(config)
     if trainer_config.policy_type == "interaction_state":
@@ -34,7 +47,7 @@ def build_policy(config: PolicyConfig, trainer_config: TrainerConfig) -> nn.Modu
 
 
 def policy_supports_planning(policy_type: str) -> bool:
-    return policy_type in {"reveal_state", "interaction_state"}
+    return policy_type in {"reveal_state", "interaction_state", "elastic_reveal"}
 
 
 def planner_enabled(trainer_config: TrainerConfig, during_eval: bool) -> bool:
@@ -93,6 +106,22 @@ class BimanualTrainer:
                 forward_kwargs["support_mode_conditioning"] = self.config.support_mode_conditioning
                 if "candidate_action_chunks" in batch:
                     forward_kwargs["candidate_chunks_override"] = batch["candidate_action_chunks"]
+            if self.config.policy_type == "elastic_reveal":
+                forward_kwargs["depths"] = batch.get("depths")
+                forward_kwargs["depth_valid"] = batch.get("depth_valid")
+                forward_kwargs["camera_intrinsics"] = batch.get("camera_intrinsics")
+                forward_kwargs["camera_extrinsics"] = batch.get("camera_extrinsics")
+                forward_kwargs["history_depths"] = batch.get("history_depths")
+                forward_kwargs["history_depth_valid"] = batch.get("history_depth_valid")
+                forward_kwargs["use_depth"] = self.config.use_depth
+                forward_kwargs["use_world_model"] = self.config.use_world_model
+                forward_kwargs["use_planner"] = planner_enabled(self.config, during_eval=False)
+                forward_kwargs["use_role_tokens"] = self.config.use_role_tokens
+                forward_kwargs["history_steps_override"] = self.config.history_steps_override
+                forward_kwargs["compute_equivariance_probe"] = self.config.compute_equivariance_probe
+            elif self.config.policy_type == "interaction_state":
+                forward_kwargs["use_role_tokens"] = self.config.use_role_tokens
+                forward_kwargs["history_steps_override"] = self.config.history_steps_override
             model_output = self.model(**forward_kwargs)
             losses = compute_total_loss(model_output, batch, weights=loss_weights)
         losses["total"].backward()