diff --git a/.gitignore b/.gitignore
index 126722d..6345b91 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,6 +23,7 @@ wandb/
runs/
outputs/
checkpoints/
+/packages/
*.ckpt
*.pt
*.pth
diff --git a/README.md b/README.md
index 33a1e34..4aee340 100644
--- a/README.md
+++ b/README.md
@@ -119,14 +119,16 @@ setting.
| Method | Clean ↑ | Random ↑ | Total ↑ |
|---|---:|---:|---:|
-| π₀ | 65.92 | 58.40 | 62.2 |
-| π₀.₅ | 82.74 | 76.76 | 79.8 |
-| Motus | **88.66** | 87.02 | **87.8** |
-| Motus from WAN2.2 | 77.56 | 77.00 | 77.3 |
-| FastWAM-Joint | 87.8 | 87.32 | 87.56 |
-| StarWAM-Joint | 84.8 | 86.0 | 85.4 |
-| StarWAM-CD | 79.0 | 79.2 | 79.1 |
-| Streaming-WAM (Ours) | 87.2 | **88.8** | 87.6 |
+| π₀ | 65.92 | 58.40 | 62.20 |
+| π₀.₅ | 82.74 | 76.76 | 79.80 |
+| Motus | 88.66 | 87.02 | 87.80 |
+| Motus from WAN2.2 | 77.56 | 77.00 | 77.30 |
+| FastWAM-Joint | 86.40 | 87.60 | 87.00 |
+| FastWAM-Joint-CD | 86.20 | 85.80 | 86.00 |
+| Streaming-WAM (FastWAM, Ours) | **90.40** | **90.80** | **90.60** |
+| StarWAM-Joint | 87.80 | 84.60 | 86.20 |
+| StarWAM-CD | 85.40 | 86.20 | 85.80 |
+| Streaming-WAM (StarWAM, Ours) | 86.60 | 85.80 | 86.20 |
#### RoboCasa
@@ -164,14 +166,17 @@ Task success alone does not characterize runtime efficiency. We therefore report
| LIBERO | Streaming-WAM | 41.0 ms | 5.36 s Long / 3.15 s Short |
| LIBERO | Streaming-WAM w/o Action Conditioning | 35.1 ms | 5.20 s Long / 2.92 s Short |
| LIBERO | Streaming-WAM w/o Slot Encoder | 36.3 ms | 5.31 s Long / 3.01 s Short |
-| RoboTwin 2.0 | StarWAM-Joint | 190.17 ms | 110.22 s |
-| RoboTwin 2.0 | StarWAM-CD | 81.21 ms | 102.59 s |
-| RoboTwin 2.0 | Streaming-WAM | 47.09 ms | 77.48 s |
+| RoboTwin 2.0 | FastWAM-Joint | 652.1 ms | 24.26 s |
+| RoboTwin 2.0 | FastWAM-Joint-CD | 165.2 ms | 18.63 s |
+| RoboTwin 2.0 | Streaming-WAM (FastWAM) | 54.4 ms | 20.14 s |
+| RoboTwin 2.0 | StarWAM-Joint | 196.5 ms | 25.76 s |
+| RoboTwin 2.0 | StarWAM-CD | 83.1 ms | 26.23 s |
+| RoboTwin 2.0 | Streaming-WAM (StarWAM) | 36.6 ms | 24.44 s |
| RoboCasa | X-WAM | 374.07 ms | 17.36 s |
| RoboCasa | X-WAM-CD | 134.37 ms | 13.04 s |
| RoboCasa | Streaming-WAM | 115.98 ms | 9.49 s |
-Across all three benchmarks, Streaming-WAM reduces both runtime measures while maintaining comparable task success. On LIBERO, it achieves a 12.0× Chunk Time speedup over FastWAM and Episode Time speedups of 3.0× and 2.6× on Long and Short tasks, respectively, with 98.20% average success. On RoboTwin 2.0, relative to StarWAM-Joint, Chunk Time falls from 190.17 ms to 47.09 ms and Episode Time from 110.22 s to 77.48 s, while overall success increases from 85.4 to 87.6. On RoboCasa, relative to X-WAM, Streaming-WAM achieves a 3.2× Chunk Time speedup and a 1.8× Episode Time speedup, with comparable average success (75.35% versus 75.42%).
+Across all three benchmarks, Streaming-WAM reduces Chunk Time while maintaining competitive task success. On LIBERO, it achieves a 12.0× Chunk Time speedup over FastWAM and Episode Time speedups of 3.0× and 2.6× on Long and Short tasks, respectively, with 98.20% average success. On RoboTwin 2.0, the FastWAM version reduces Chunk Time from 652.1 ms to 54.4 ms and Episode Time from 24.26 s to 20.14 s, while success rises from 87.00 to 90.60. The StarWAM version reduces Chunk Time from 196.5 ms to 36.6 ms and Episode Time from 25.76 s to 24.44 s, with the same 86.20 overall success. On RoboCasa, relative to X-WAM, Streaming-WAM achieves a 3.2× Chunk Time speedup and a 1.8× Episode Time speedup, with comparable average success (75.35% versus 75.42%).
## Runtime layout
diff --git a/docs/assets/streaming-wam-chunk-time.png b/docs/assets/streaming-wam-chunk-time.png
index d249450..f283c20 100644
Binary files a/docs/assets/streaming-wam-chunk-time.png and b/docs/assets/streaming-wam-chunk-time.png differ
diff --git a/docs/assets/streaming-wam-episode-time.png b/docs/assets/streaming-wam-episode-time.png
index 8766ce7..d3c611c 100644
Binary files a/docs/assets/streaming-wam-episode-time.png and b/docs/assets/streaming-wam-episode-time.png differ
diff --git a/docs/generate_latency_figure.py b/docs/generate_latency_figure.py
index 031bd3e..b4c4c66 100644
--- a/docs/generate_latency_figure.py
+++ b/docs/generate_latency_figure.py
@@ -41,8 +41,8 @@
LIBERO_SHORT = (8.25, 3.74, 3.20, 3.15, 2.92, 3.01)
ROBOTWIN_METHODS = ("StarWAM\nJoint", "StarWAM\nCD", "Streaming-\nWAM")
-ROBOTWIN_CHUNK = (190.17, 81.21, 47.09)
-ROBOTWIN_EPISODE = (110.22, 102.59, 77.48)
+ROBOTWIN_CHUNK = (196.5, 83.1, 36.6)
+ROBOTWIN_EPISODE = (25.76, 26.23, 24.44)
ROBOCASA_METHODS = ("X-WAM", "X-WAM\nCD", "Streaming-\nWAM")
ROBOCASA_CHUNK = (374.07, 134.37, 115.98)
ROBOCASA_EPISODE = (17.36, 13.04, 9.49)
@@ -240,7 +240,7 @@ def render_episode_time(output_path: Path) -> None:
_draw_libero_episode(axes[0])
_draw_three_method_panel(
axes[1], title="RoboTwin 2.0", ylabel="Seconds", methods=ROBOTWIN_METHODS,
- values=ROBOTWIN_EPISODE, ceiling=125,
+ values=ROBOTWIN_EPISODE, ceiling=30,
)
_draw_three_method_panel(
axes[2], title="RoboCasa", ylabel="Seconds", methods=ROBOCASA_METHODS,
diff --git a/docs/index.html b/docs/index.html
index 876208f..acdfa0f 100644
--- a/docs/index.html
+++ b/docs/index.html
@@ -302,14 +302,16 @@
RoboTwin 2.0
RoboTwin 2.0 clean and randomized results
Method Clean ↑ Random ↑ Total ↑
- π₀ 65.92 58.40 62.2
- π₀.₅ 82.74 76.76 79.8
- Motus 88.66 87.02 87.8
- Motus from WAN2.2 77.56 77.00 77.3
- FastWAM-Joint 87.8 87.32 87.56
- StarWAM-Joint 84.8 86.0 85.4
- StarWAM-CD 79.0 79.2 79.1
- Streaming-WAM (Ours) 87.2 88.8 87.6
+ π₀ 65.92 58.40 62.20
+ π₀.₅ 82.74 76.76 79.80
+ Motus 88.66 87.02 87.80
+ Motus from WAN2.2 77.56 77.00 77.30
+ FastWAM-Joint 86.40 87.60 87.00
+ FastWAM-Joint-CD 86.20 85.80 86.00
+ Streaming-WAM (FastWAM, Ours) 90.40 90.80 90.60
+ StarWAM-Joint 87.80 84.60 86.20
+ StarWAM-CD 85.40 86.20 85.80
+ Streaming-WAM (StarWAM, Ours) 86.60 85.80 86.20
@@ -357,7 +359,7 @@ Real robot evaluation
Inference efficiency. Task success alone does not characterize runtime efficiency. We therefore report Chunk Time, the latency required to prepare the next action chunk, and Episode Time, the duration of a complete rollout, including inference, execution, and replanning.
-
Across all three benchmarks, Streaming-WAM reduces both runtime measures while maintaining comparable task success. On LIBERO, it achieves a 12.0× Chunk Time speedup over FastWAM and Episode Time speedups of 3.0× and 2.6× on Long and Short tasks, respectively, with 98.20% average success. On RoboTwin 2.0, relative to StarWAM-Joint, Chunk Time falls from 190.17 ms to 47.09 ms and Episode Time from 110.22 s to 77.48 s, while overall success increases from 85.4 to 87.6. On RoboCasa, relative to X-WAM, Streaming-WAM achieves a 3.2× Chunk Time speedup and a 1.8× Episode Time speedup, with comparable average success (75.35% versus 75.42%).
+
Across all three benchmarks, Streaming-WAM reduces Chunk Time while maintaining competitive task success. On LIBERO, it achieves a 12.0× Chunk Time speedup over FastWAM and Episode Time speedups of 3.0× and 2.6× on Long and Short tasks, respectively, with 98.20% average success. On RoboTwin 2.0, the FastWAM version reduces Chunk Time from 652.1 ms to 54.4 ms and Episode Time from 24.26 s to 20.14 s, while success rises from 87.00 to 90.60. The StarWAM version reduces Chunk Time from 196.5 ms to 36.6 ms and Episode Time from 25.76 s to 24.44 s, with the same 86.20 overall success. On RoboCasa, relative to X-WAM, Streaming-WAM achieves a 3.2× Chunk Time speedup and a 1.8× Episode Time speedup, with comparable average success (75.35% versus 75.42%).
@@ -385,9 +387,12 @@
Real robot evaluation
LIBERO Streaming-WAM 41.0 ms 5.36 s Long / 3.15 s Short
LIBERO Streaming-WAM w/o Action Conditioning 35.1 ms 5.20 s Long / 2.92 s Short
LIBERO Streaming-WAM w/o Slot Encoder 36.3 ms 5.31 s Long / 3.01 s Short
-
RoboTwin 2.0 StarWAM-Joint 190.17 ms 110.22 s
-
RoboTwin 2.0 StarWAM-CD 81.21 ms 102.59 s
-
RoboTwin 2.0 Streaming-WAM 47.09 ms 77.48 s
+
RoboTwin 2.0 FastWAM-Joint 652.1 ms 24.26 s
+
RoboTwin 2.0 FastWAM-Joint-CD 165.2 ms 18.63 s
+
RoboTwin 2.0 Streaming-WAM (FastWAM) 54.4 ms 20.14 s
+
RoboTwin 2.0 StarWAM-Joint 196.5 ms 25.76 s
+
RoboTwin 2.0 StarWAM-CD 83.1 ms 26.23 s
+
RoboTwin 2.0 Streaming-WAM (StarWAM) 36.6 ms 24.44 s
RoboCasa X-WAM 374.07 ms 17.36 s
RoboCasa X-WAM-CD 134.37 ms 13.04 s
RoboCasa Streaming-WAM 115.98 ms 9.49 s
diff --git a/examples/robotwin/RoboTwin.md b/examples/robotwin/RoboTwin.md
index f8e57e1..a79fc2b 100644
--- a/examples/robotwin/RoboTwin.md
+++ b/examples/robotwin/RoboTwin.md
@@ -261,7 +261,42 @@ python script/eval_policy.py \
For Shared-DiT, also override `config_path` to the Shared-DiT recipe, use
checkpoint-45000, and set both inference-step values to 16.
-### 8.3 Results
+### 8.3 Canonical cross-family benchmark
+
+Use the same public harness for StarWAM, FastWAM, and Streaming-WAM checkpoints.
+The wrapper defaults to 100 trials per task-setting and `replan_steps=16`; every
+path and runtime remains overridable through environment variables.
+
+```bash
+MODE=baseline \
+CHECKPOINT_FORMAT=fastwam \
+CKPT=/path/to/fastwam_joint.pt \
+CONFIG=examples/robotwin/configs/recipes/streamingwam_robotwin_mot_wan22_5b.yaml \
+STATS_PATH=/path/to/action_stats.json \
+BACKBONE_PATH=/path/to/Wan2.2-TI2V-5B \
+ROBOTWIN_HOME=/path/to/RoboTwin \
+INFERENCE_PYTHON=/path/to/streamingwam-env/bin/python \
+SIMULATOR_PYTHON=/path/to/robotwin-env/bin/python \
+bash examples/robotwin/scripts/run_streamingwam_robotwin_benchmark.sh
+```
+
+The released defaults are selected from `CHECKPOINT_FORMAT` and `MODE`:
+
+| Checkpoint family | Mode | Video steps | Action steps |
+|---|---|---:|---:|
+| StarWAM | baseline | 4 | 4 |
+| StarWAM | cd | 1 | 1 |
+| StarWAM | ac-stream | 1 | 1 |
+| FastWAM | baseline | 10 | 10 |
+| FastWAM | cd | 1 | 2 |
+| Streaming-WAM/FastWAM Stage2 | ac-stream | 1 | 2 |
+
+Set `NUM_INFERENCE_STEPS` and `ACTION_NUM_INFERENCE_STEPS` to override these
+defaults for another checkpoint. Use `AC_STREAM_BACKEND=eager` for the eager
+ablation; accelerated AC-Stream is the default. `MODEL_SEED` controls model
+noise and `EPISODE_SEED` independently controls the RoboTwin scene chain.
+
+### 8.4 Results
RoboTwin writes one result per (task, config) under
`RoboTwin/eval_result/
/////_result.txt`,
diff --git a/examples/robotwin/client_policy.py b/examples/robotwin/client_policy.py
index 3cab78a..685f093 100644
--- a/examples/robotwin/client_policy.py
+++ b/examples/robotwin/client_policy.py
@@ -65,6 +65,42 @@ def _is_none_like(value: Any) -> bool:
return False
+def _take_environment_action(
+ task_env: Any,
+ action: np.ndarray,
+ *,
+ defer_render: bool,
+) -> tuple[int, int]:
+ """Execute one action while suppressing redundant renders during D8 overlap."""
+
+ original_update_render = None
+ original_get_obs = None
+ if defer_render:
+ original_update_render = task_env._update_render
+ original_get_obs = task_env.get_obs
+
+ def skip_update_render(*_args, **_kwargs):
+ return None
+
+ def get_obs_with_render(*args, **kwargs):
+ task_env._update_render = original_update_render
+ try:
+ return original_get_obs(*args, **kwargs)
+ finally:
+ task_env._update_render = skip_update_render
+
+ task_env._update_render = skip_update_render
+ task_env.get_obs = get_obs_with_render
+ started_ns = time.perf_counter_ns()
+ try:
+ task_env.take_action(action, action_type="qpos")
+ finally:
+ if original_update_render is not None:
+ task_env._update_render = original_update_render
+ task_env.get_obs = original_get_obs
+ return started_ns, time.perf_counter_ns()
+
+
class RemoteStreamingWAMModel:
"""Talks to the Streaming-WAM inference server; manages the replan queue locally."""
@@ -94,9 +130,16 @@ def __init__(
self._timing_records: list[dict[str, Any]] = []
self._episode_records: list[dict[str, Any]] = []
self._episode_started_ns: int | None = None
+ self._terminal_action_completion_ns: int | None = None
+ self._timing_hook_active = False
+ self._active_timing_metadata: dict[str, Any] = {}
self._episode_index = 0
self._replan_index = 0
+ @property
+ def needs_prewarm(self) -> bool:
+ return self.prewarm and not self._prewarmed
+
def _connect(self, timeout: float) -> socket.socket:
deadline = time.time() + timeout
last_err: Optional[Exception] = None
@@ -252,15 +295,24 @@ def _prewarm_once(self, task_env: Any, observation: Dict[str, Any] | None) -> No
self._prewarmed = True
self._replan_index = 0
- def _start_episode(self) -> None:
+ def _start_episode(self, *, started_ns: int | None = None) -> None:
if self._episode_started_ns is None:
- self._episode_started_ns = time.perf_counter_ns()
+ self._episode_started_ns = (
+ time.perf_counter_ns() if started_ns is None else int(started_ns)
+ )
+ self._terminal_action_completion_ns = None
self._replan_index = 0
- def finish_episode(self, success: bool | None = None) -> None:
+ def finish_episode(
+ self,
+ success: bool | None = None,
+ *,
+ ended_ns: int | None = None,
+ ) -> None:
if self._episode_started_ns is None:
return
- elapsed = (time.perf_counter_ns() - self._episode_started_ns) / 1e9
+ completed_ns = time.perf_counter_ns() if ended_ns is None else int(ended_ns)
+ elapsed = (completed_ns - self._episode_started_ns) / 1e9
self._episode_records.append({
"record_type": "episode",
"task": self.current_task,
@@ -268,8 +320,12 @@ def finish_episode(self, success: bool | None = None) -> None:
"episode": self._episode_index,
"success": success,
"total_time_s": elapsed,
+ **self._active_timing_metadata,
})
self._episode_started_ns = None
+ self._terminal_action_completion_ns = None
+ self._timing_hook_active = False
+ self._active_timing_metadata = {}
# Drain a terminal episode's final D8 request before the worker changes
# task identity. The elapsed value above excludes this post-terminal
# cleanup, while the launched chunk remains represented in timing.
@@ -282,16 +338,24 @@ def step(self, task_env: Any, observation: Optional[Dict[str, Any]]) -> None:
self._prewarm_once(task_env, observation)
self._start_episode()
if self.inference_mode == "ac-stream":
- if observation is None:
- raise ValueError("AC-Stream requires an observation at every action step")
- snapshot = self._snapshot(task_env, observation)
if self._ac_stream_controller is None:
+ if observation is None:
+ raise ValueError("AC-Stream D0 requires an observation")
+ snapshot = self._snapshot(task_env, observation)
self._ac_stream_controller = ACStreamController(self._predict_remote)
self._ac_stream_controller.start_episode(snapshot)
+ elif self._ac_stream_controller.observation_required:
+ if observation is None:
+ raise ValueError("AC-Stream D8 launch requires an observation")
+ snapshot = self._snapshot(task_env, observation)
+ else:
+ snapshot = None
action = self._ac_stream_controller.next_action(snapshot)
- started_ns = time.perf_counter_ns()
- task_env.take_action(action, action_type="qpos")
- completed_ns = time.perf_counter_ns()
+ started_ns, completed_ns = _take_environment_action(
+ task_env,
+ action,
+ defer_render=self._ac_stream_controller.inference_in_flight,
+ )
self._ac_stream_controller.mark_action_executed(
started_ns=started_ns,
completed_ns=completed_ns,
@@ -311,11 +375,71 @@ def step(self, task_env: Any, observation: Optional[Dict[str, Any]]) -> None:
for i in range(min(self.replan_steps, chunk.shape[0])):
self.pending_actions.append(np.asarray(chunk[i], dtype=np.float32))
if self.pending_actions:
- task_env.take_action(self.pending_actions.popleft(), action_type="qpos")
- if bool(getattr(task_env, "eval_success", False)):
- self.finish_episode(success=True)
- elif int(getattr(task_env, "take_action_cnt", 0)) >= int(getattr(task_env, "step_lim", 1 << 30)):
- self.finish_episode(success=False)
+ _, completed_ns = _take_environment_action(
+ task_env,
+ self.pending_actions.popleft(),
+ defer_render=False,
+ )
+ succeeded = bool(getattr(task_env, "eval_success", False))
+ exhausted = int(getattr(task_env, "take_action_cnt", 0)) >= int(
+ getattr(task_env, "step_lim", 1 << 30)
+ )
+ if succeeded or exhausted:
+ self._terminal_action_completion_ns = completed_ns
+ if not self._timing_hook_active:
+ self.finish_episode(
+ success=succeeded,
+ ended_ns=self._terminal_action_completion_ns,
+ )
+
+ def should_request_observation(self) -> bool:
+ if self.inference_mode != "ac-stream":
+ return not self.pending_actions
+ if self._ac_stream_controller is None:
+ return True
+ return self._ac_stream_controller.observation_required
+
+ def prepare_instruction(self, instruction: str) -> None:
+ self._request_id += 1
+ _send_msg(
+ self._conn,
+ {
+ "cmd": "prepare_instruction",
+ "request_id": self._request_id,
+ "instruction": str(instruction),
+ },
+ )
+ response = _recv_msg(self._conn)
+ if "error" in response:
+ raise RuntimeError(f"policy server error:\n{response['error']}")
+ if response.get("request_id") != self._request_id:
+ raise RuntimeError(
+ f"policy response ID mismatch: expected {self._request_id}, "
+ f"got {response.get('request_id')}"
+ )
+
+ def prewarm_model(self, task_env: Any, observation: Dict[str, Any]) -> None:
+ self._prewarm_once(task_env, observation)
+
+ def begin_timing_trajectory(self, metadata: dict[str, Any]) -> None:
+ if self._episode_started_ns is not None:
+ raise RuntimeError("a timed RoboTwin trajectory is already active")
+ self._timing_hook_active = True
+ self._active_timing_metadata = dict(metadata)
+ self._start_episode(started_ns=time.perf_counter_ns())
+
+ def end_timing_trajectory(
+ self,
+ success: bool,
+ metadata: dict[str, Any],
+ ) -> None:
+ if not self._timing_hook_active or self._episode_started_ns is None:
+ raise RuntimeError("no timed RoboTwin trajectory is active")
+ ended_ns = self._terminal_action_completion_ns
+ if ended_ns is None:
+ ended_ns = time.perf_counter_ns()
+ self._active_timing_metadata.update(metadata)
+ self.finish_episode(success=bool(success), ended_ns=ended_ns)
def reset(self) -> None:
self.finish_episode(success=None)
@@ -372,5 +496,32 @@ def reset_model(model: RemoteStreamingWAMModel) -> None:
model.reset()
+def prepare_instruction(task_env: Any, model: RemoteStreamingWAMModel) -> None:
+ model.prepare_instruction(task_env.get_instruction())
+
+
+def prewarm_model(
+ task_env: Any,
+ model: RemoteStreamingWAMModel,
+ observation: Dict[str, Any],
+) -> None:
+ model.prewarm_model(task_env, observation)
+
+
+def begin_timing_trajectory(
+ model: RemoteStreamingWAMModel,
+ metadata: dict[str, Any],
+) -> None:
+ model.begin_timing_trajectory(metadata)
+
+
+def end_timing_trajectory(
+ model: RemoteStreamingWAMModel,
+ success: bool,
+ metadata: dict[str, Any],
+) -> None:
+ model.end_timing_trajectory(success, metadata)
+
+
def _finish_episode(model: RemoteStreamingWAMModel) -> None:
model.finish_episode(success=None)
diff --git a/examples/robotwin/deploy_policy.py b/examples/robotwin/deploy_policy.py
index d409830..0d87982 100644
--- a/examples/robotwin/deploy_policy.py
+++ b/examples/robotwin/deploy_policy.py
@@ -68,7 +68,31 @@ def _load_adapter(mode: str):
def get_model(usr_args: Dict[str, Any]):
mode = _infer_mode(usr_args)
- return _load_adapter(mode).get_model(usr_args)
+ adapter = _load_adapter(mode)
+ model = adapter.get_model(usr_args)
+ # Keep the selected adapter on the model so every public deploy-policy
+ # hook is forwarded through the same implementation. RoboTwin imports
+ # this module as the stable entry point; evaluators must not depend on a
+ # concrete client/local adapter module.
+ model._streamingwam_robotwin_adapter = adapter
+ return model
+
+
+def _adapter(model: Any):
+ adapter = getattr(model, "_streamingwam_robotwin_adapter", None)
+ if adapter is None:
+ raise RuntimeError("model was not created by deploy_policy.get_model")
+ return adapter
+
+
+def _required_hook(model: Any, name: str, *args: Any) -> Any:
+ hook = getattr(_adapter(model), name, None)
+ if not callable(hook):
+ raise AttributeError(
+ f"RoboTwin policy adapter {_adapter(model).__name__!r} "
+ f"does not implement required hook {name!r}"
+ )
+ return hook(*args)
def eval(TASK_ENV: Any, model: Any, observation: Optional[Dict[str, Any]]) -> None:
@@ -79,6 +103,30 @@ def reset_model(model: Any) -> None:
model.reset()
+def prepare_instruction(task_env: Any, model: Any) -> None:
+ _required_hook(model, "prepare_instruction", task_env, model)
+
+
+def prewarm_model(
+ task_env: Any,
+ model: Any,
+ observation: Dict[str, Any],
+) -> None:
+ _required_hook(model, "prewarm_model", task_env, model, observation)
+
+
+def begin_timing_trajectory(model: Any, metadata: Dict[str, Any]) -> None:
+ _required_hook(model, "begin_timing_trajectory", model, metadata)
+
+
+def end_timing_trajectory(
+ model: Any,
+ success: bool,
+ metadata: Dict[str, Any],
+) -> None:
+ _required_hook(model, "end_timing_trajectory", model, success, metadata)
+
+
def _finish_episode(model: Any) -> None:
finish = getattr(model, "finish_episode", None)
if callable(finish):
diff --git a/examples/robotwin/evaluator.py b/examples/robotwin/evaluator.py
new file mode 100644
index 0000000..3f43799
--- /dev/null
+++ b/examples/robotwin/evaluator.py
@@ -0,0 +1,136 @@
+"""Canonical one-episode RoboTwin evaluation used by the public benchmark."""
+
+from __future__ import annotations
+
+import importlib
+import traceback
+from typing import Any
+
+import numpy as np
+
+
+def evaluate_one_episode(
+ *,
+ official: Any,
+ task_name: str,
+ task_env: Any,
+ settings: dict[str, Any],
+ model: Any,
+ start_seed: int,
+ trial: int,
+ instruction_type: str = "unseen",
+ policy_module: Any | None = None,
+) -> tuple[int, int]:
+ """Run one accepted RoboTwin episode with canonical trajectory timing.
+
+ Scene construction, expert validation, instruction preparation, and model
+ warmup are outside the timer. Timing begins immediately before the first
+ policy observation and ends when the terminal action call returns.
+ """
+
+ if getattr(task_env, "eval_video_path", None) is not None:
+ raise RuntimeError("benchmark timing requires eval_video_log=false")
+
+ if policy_module is None:
+ policy_module = importlib.import_module(str(settings["policy_name"]))
+ reset_model = getattr(policy_module, "reset_model")
+ prepare_instruction = getattr(policy_module, "prepare_instruction")
+ prewarm_model = getattr(policy_module, "prewarm_model")
+ begin_timing = getattr(policy_module, "begin_timing_trajectory")
+ end_timing = getattr(policy_module, "end_timing_trajectory")
+
+ task_env.suc = 0
+ task_env.test_num = 0
+ now_seed = int(start_seed)
+ settings["eval_mode"] = True
+
+ while True:
+ render_freq = settings["render_freq"]
+ settings["render_freq"] = 0
+ try:
+ task_env.setup_demo(
+ now_ep_num=0,
+ seed=now_seed,
+ is_test=True,
+ **settings,
+ )
+ episode_info = task_env.play_once()
+ task_env.close_env()
+ except official.UnStableError:
+ task_env.close_env()
+ now_seed += 1
+ settings["render_freq"] = render_freq
+ continue
+ except Exception:
+ traceback.print_exc()
+ task_env.close_env()
+ now_seed += 1
+ settings["render_freq"] = render_freq
+ continue
+
+ if not (task_env.plan_success and task_env.check_success()):
+ now_seed += 1
+ settings["render_freq"] = render_freq
+ continue
+
+ settings["render_freq"] = render_freq
+ try:
+ task_env.setup_demo(
+ now_ep_num=0,
+ seed=now_seed,
+ is_test=True,
+ **settings,
+ )
+ except official.UnStableError:
+ task_env.close_env()
+ now_seed += 1
+ continue
+
+ descriptions = official.generate_episode_descriptions(
+ settings["task_name"],
+ [episode_info["info"]],
+ 1,
+ )
+ instruction = np.random.choice(descriptions[0][instruction_type])
+ task_env.set_instruction(instruction=instruction)
+
+ prepare_instruction(task_env, model)
+ if bool(getattr(model, "needs_prewarm", False)):
+ prewarm_model(task_env, model, task_env.get_obs())
+ reset_model(model)
+
+ metadata = {
+ "task": str(task_name),
+ "config": str(settings["task_config"]),
+ "trial": int(trial),
+ "accepted_seed": int(now_seed),
+ }
+ begin_timing(model, metadata)
+ success = False
+ while task_env.take_action_cnt < task_env.step_lim:
+ observation = (
+ task_env.get_obs()
+ if bool(model.should_request_observation())
+ else None
+ )
+ policy_module.eval(task_env, model, observation)
+ if task_env.eval_success:
+ success = True
+ break
+ end_timing(
+ model,
+ success,
+ {
+ **metadata,
+ "environment_action_count": int(task_env.take_action_cnt),
+ "environment_step_limit": int(task_env.step_lim),
+ },
+ )
+
+ if success:
+ task_env.suc = 1
+ task_env.test_num = 1
+ task_env.close_env()
+ if task_env.render_freq:
+ task_env.viewer.close()
+ return now_seed + 1, int(success)
diff --git a/examples/robotwin/multigpu_rollout.py b/examples/robotwin/multigpu_rollout.py
index adc7528..c25c2f3 100644
--- a/examples/robotwin/multigpu_rollout.py
+++ b/examples/robotwin/multigpu_rollout.py
@@ -27,6 +27,29 @@ def _csv(value: str) -> list[str]:
return [item.strip() for item in value.split(",") if item.strip()]
+def resolve_sampling_steps(
+ checkpoint_format: str,
+ inference_mode: str,
+ video_steps: int | None,
+ action_steps: int | None,
+) -> tuple[int, int]:
+ """Resolve the released RoboTwin recipes while preserving CLI overrides."""
+
+ if video_steps is None:
+ if checkpoint_format == "fastwam" and inference_mode == "baseline":
+ video_steps = 10
+ elif checkpoint_format == "starwam" and inference_mode == "baseline":
+ video_steps = 4
+ else:
+ video_steps = 1
+ if action_steps is None:
+ if checkpoint_format in {"fastwam", "streamingwam"} and inference_mode != "baseline":
+ action_steps = 2
+ else:
+ action_steps = video_steps
+ return int(video_steps), int(action_steps)
+
+
def _build_arg_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Streaming-WAM RoboTwin multi-GPU evaluation")
parser.add_argument("--config", required=True)
@@ -39,14 +62,31 @@ def _build_arg_parser() -> argparse.ArgumentParser:
parser.add_argument("--inference-python", required=True)
parser.add_argument("--simulator-python", required=True)
parser.add_argument("--inference-mode", choices=("baseline", "cd", "ac-stream"), required=True)
+ parser.add_argument("--num-inference-steps", type=int, default=None)
+ parser.add_argument("--action-num-inference-steps", type=int, default=None)
backend = parser.add_mutually_exclusive_group()
backend.add_argument("--ac-stream-accelerated", action="store_true")
backend.add_argument("--ac-stream-eager", action="store_true")
parser.add_argument("--gpu-ids", default=os.environ.get("GPU_IDS", "0,1,2,3"))
+ parser.add_argument("--inference-gpu-ids", default=None)
+ parser.add_argument("--simulator-gpu-ids", default=None)
+ parser.add_argument("--require-gpu-isolation", action="store_true")
parser.add_argument("--tasks", default=None)
parser.add_argument("--configs", default="demo_clean,demo_randomized")
parser.add_argument("--num-trials", type=int, default=1)
- parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--replan-steps", type=int, default=16)
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="model inference/noise seed; keep fixed across benchmark retries",
+ )
+ parser.add_argument(
+ "--episode-seed",
+ type=int,
+ default=None,
+ help="RoboTwin scene seed chain; defaults to --seed for compatibility",
+ )
parser.add_argument("--port-base", type=int, default=18765)
parser.add_argument("--job-timeout-seconds", type=float, default=1200.0)
parser.add_argument("--output-dir", default=None)
@@ -63,13 +103,19 @@ def build_server_command(args: argparse.Namespace, *, port: int) -> list[str]:
accelerated=accelerated,
eager=args.ac_stream_eager,
)
+ video_steps, action_steps = resolve_sampling_steps(
+ args.checkpoint_format,
+ args.inference_mode,
+ args.num_inference_steps,
+ args.action_num_inference_steps,
+ )
command = [
args.inference_python, "-m", "examples.robotwin.policy_server",
"--config", args.config, "--checkpoint", args.checkpoint,
"--checkpoint-format", args.checkpoint_format,
"--inference-mode", args.inference_mode,
- "--num-inference-steps", "4" if args.inference_mode == "baseline" else "1",
- "--action-num-inference-steps", "4" if args.inference_mode == "baseline" else "1",
+ "--num-inference-steps", str(video_steps),
+ "--action-num-inference-steps", str(action_steps),
"--seed", str(args.seed),
"--device", "cuda:0", "--host", "127.0.0.1", "--port", str(port),
"--override",
@@ -150,7 +196,7 @@ def _select_pending_jobs(jobs, output_dir: Path) -> tuple[list, list[dict]]:
except (OSError, json.JSONDecodeError):
pending.append(job)
continue
- if result.get("status") in {"completed", "skipped_timeout", "infrastructure_error"}:
+ if result.get("status") == "completed":
existing.append(result)
else:
pending.append(job)
@@ -197,6 +243,19 @@ def _atomic_json(path: Path, payload: dict) -> None:
temporary.replace(path)
+def _ensure_run_config(path: Path, current: dict) -> None:
+ """Prevent a resumed output directory from mixing benchmark protocols."""
+
+ if path.is_file():
+ previous = json.loads(path.read_text(encoding="utf-8"))
+ if previous != current:
+ raise RuntimeError(
+ f"run configuration mismatch in {path}; use a new --output-dir"
+ )
+ return
+ _atomic_json(path, current)
+
+
def _read_phase(path: Path) -> str:
try:
return str(json.loads(path.read_text(encoding="utf-8")).get("phase", "unknown"))
@@ -298,18 +357,19 @@ def launch(item: dict, job) -> None:
status_output.unlink(missing_ok=True)
_atomic_json(job_file, job.__dict__)
log = open(job_dir / "worker.log", "w", encoding="utf-8")
- replan_steps = 16 if args.inference_mode == "ac-stream" else 32
command = [
args.simulator_python, "-m", "examples.robotwin.robotwin_worker",
"--gpu-id", gpu, "--robotwin-home", args.robotwin_home,
"--policy-dir", str(_ROOT / "examples" / "robotwin"),
"--server-port", str(item["port"]),
"--inference-mode", args.inference_mode,
- "--replan-steps", str(replan_steps),
+ "--replan-steps", str(args.replan_steps),
"--job-file", str(job_file),
"--output", str(worker_output),
"--status-output", str(status_output),
- "--seed", str(args.seed),
+ "--seed", str(
+ args.seed if args.episode_seed is None else args.episode_seed
+ ),
"--prewarm" if not warmed[gpu] else "--no-prewarm",
]
try:
@@ -454,7 +514,8 @@ def _terminate_processes(processes: list) -> None:
def _start_server_group(
args: argparse.Namespace,
- gpu_ids: list[str],
+ inference_gpu_ids: list[str],
+ simulator_gpu_ids: list[str],
output_dir: Path,
*,
popen=subprocess.Popen,
@@ -464,21 +525,31 @@ def _start_server_group(
group: list[dict] = []
try:
- for worker_index, gpu in enumerate(gpu_ids):
- worker_dir = output_dir / f"worker_gpu{gpu}"
+ for worker_index, (inference_gpu, simulator_gpu) in enumerate(
+ zip(inference_gpu_ids, simulator_gpu_ids)
+ ):
+ worker_dir = output_dir / (
+ f"pair_{worker_index}_inference_gpu{inference_gpu}_"
+ f"simulator_gpu{simulator_gpu}"
+ )
worker_dir.mkdir(parents=True, exist_ok=True)
port = args.port_base + worker_index
- environment = os.environ.copy()
- environment["CUDA_VISIBLE_DEVICES"] = gpu
- environment["PYTHONPATH"] = (
- str(_ROOT) + os.pathsep + environment.get("PYTHONPATH", "")
+ server_environment = os.environ.copy()
+ server_environment["CUDA_VISIBLE_DEVICES"] = inference_gpu
+ server_environment["PYTHONPATH"] = (
+ str(_ROOT) + os.pathsep + server_environment.get("PYTHONPATH", "")
+ )
+ simulator_environment = os.environ.copy()
+ simulator_environment["CUDA_VISIBLE_DEVICES"] = simulator_gpu
+ simulator_environment["PYTHONPATH"] = (
+ str(_ROOT) + os.pathsep + simulator_environment.get("PYTHONPATH", "")
)
log = open(worker_dir / "server.log", "w", encoding="utf-8")
try:
process = popen(
build_server_command(args, port=port),
cwd=_ROOT,
- env=environment,
+ env=server_environment,
stdout=log,
stderr=subprocess.STDOUT,
start_new_session=True,
@@ -487,10 +558,11 @@ def _start_server_group(
log.close()
raise
group.append({
- "gpu": gpu,
+ "gpu": simulator_gpu,
+ "inference_gpu": inference_gpu,
"port": port,
"worker_dir": worker_dir,
- "environment": environment,
+ "environment": simulator_environment,
"process": process,
"log": log,
})
@@ -506,7 +578,25 @@ def _start_server_group(
def main() -> None:
args = _build_arg_parser().parse_args()
- gpu_ids = _csv(args.gpu_ids)
+ inference_gpu_ids = _csv(args.inference_gpu_ids or args.gpu_ids)
+ simulator_gpu_ids = _csv(args.simulator_gpu_ids or args.gpu_ids)
+ if len(inference_gpu_ids) != len(simulator_gpu_ids):
+ raise ValueError(
+ "--inference-gpu-ids and --simulator-gpu-ids must contain the same "
+ "number of GPUs"
+ )
+ if not inference_gpu_ids:
+ raise ValueError("At least one inference/simulator GPU pair is required")
+ overlap = sorted(set(inference_gpu_ids) & set(simulator_gpu_ids))
+ if args.require_gpu_isolation and overlap:
+ raise ValueError(
+ "canonical timing requires disjoint inference and simulator GPUs; "
+ f"overlap={overlap}"
+ )
+ if args.inference_mode == "ac-stream" and args.replan_steps != 16:
+ raise ValueError("AC-Stream has fixed H=32, s=16, d=8 geometry")
+ if not 1 <= args.replan_steps <= 32:
+ raise ValueError("--replan-steps must be in [1, 32]")
jobs = build_workload(
num_trials=args.num_trials,
tasks=_csv(args.tasks) if args.tasks else None,
@@ -515,16 +605,48 @@ def main() -> None:
stamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
output_dir = Path(args.output_dir or (_ROOT / "outputs" / f"robotwin_multigpu_{stamp}")).resolve()
output_dir.mkdir(parents=True, exist_ok=True)
+ _ensure_run_config(
+ output_dir / "run_config.json",
+ {
+ "config": str(Path(args.config).resolve()),
+ "checkpoint": str(Path(args.checkpoint).resolve()),
+ "checkpoint_format": args.checkpoint_format,
+ "stats_path": str(Path(args.stats_path).resolve()),
+ "backbone_path": str(Path(args.backbone_path).resolve()),
+ "inference_mode": args.inference_mode,
+ "num_inference_steps": args.num_inference_steps,
+ "action_num_inference_steps": args.action_num_inference_steps,
+ "replan_steps": args.replan_steps,
+ "num_trials": args.num_trials,
+ "seed": args.seed,
+ **(
+ {"episode_seed": args.episode_seed}
+ if args.episode_seed is not None
+ and args.episode_seed != args.seed
+ else {}
+ ),
+ "tasks": list(dict.fromkeys(job.task for job in jobs)),
+ "configs": list(dict.fromkeys(job.config for job in jobs)),
+ "ac_stream_accelerated": bool(args.ac_stream_accelerated),
+ "ac_stream_eager": bool(args.ac_stream_eager),
+ },
+ )
server_group: list[dict] = []
try:
pending, existing = _select_pending_jobs(jobs, output_dir)
if pending:
print(
- f"Starting {len(gpu_ids)} inference servers on GPUs "
- f"{','.join(gpu_ids)} ...",
+ f"Starting {len(inference_gpu_ids)} inference servers on GPUs "
+ f"{','.join(inference_gpu_ids)}; simulator GPUs "
+ f"{','.join(simulator_gpu_ids)} ...",
flush=True,
)
- server_group = _start_server_group(args, gpu_ids, output_dir)
+ server_group = _start_server_group(
+ args,
+ inference_gpu_ids,
+ simulator_gpu_ids,
+ output_dir,
+ )
print(
f"All {len(server_group)} inference servers are ready; "
f"dynamically scheduling {len(pending)} RoboTwin jobs ...",
@@ -576,13 +698,27 @@ def main() -> None:
})
results_path = output_dir / "results.json"
_atomic_json(results_path, {
- "mode": mode_label, "gpu_ids": gpu_ids,
+ "mode": mode_label,
+ "inference_gpu_ids": inference_gpu_ids,
+ "simulator_gpu_ids": simulator_gpu_ids,
+ "replan_steps": args.replan_steps,
+ "num_inference_steps": args.num_inference_steps,
+ "action_num_inference_steps": args.action_num_inference_steps,
"backend": backend_values,
"runtime": runtime_values[0] if runtime_values else {},
"summary": summary, "jobs": task_results, "timing": records,
})
+ _atomic_json(output_dir / "per_task_summary.json", summary["by_task"])
+ _atomic_json(output_dir / "per_setting_summary.json", summary["by_setting"])
print(format_summary(
- mode=mode_label, gpu_ids=gpu_ids, jobs=len(jobs),
+ mode=mode_label,
+ gpu_ids=[
+ f"{inference}->{simulator}"
+ for inference, simulator in zip(
+ inference_gpu_ids, simulator_gpu_ids
+ )
+ ],
+ jobs=len(jobs),
summary=summary, results_path=results_path,
completed=completed,
skipped=skipped + infrastructure_errors,
diff --git a/examples/robotwin/policy_server.py b/examples/robotwin/policy_server.py
index 8502f18..5cfb5f9 100644
--- a/examples/robotwin/policy_server.py
+++ b/examples/robotwin/policy_server.py
@@ -128,6 +128,9 @@ def handle_request(policy: StreamingWAMPolicy, req: dict) -> dict:
if cmd == "reset":
policy.reset()
return {"ok": True, "request_id": request_id}
+ if cmd == "prepare_instruction":
+ policy._encode_context(str(req["instruction"]))
+ return {"ok": True, "request_id": request_id}
if cmd != "infer":
raise ValueError(f"unsupported policy command {cmd!r}")
server_started_ns = time.perf_counter_ns()
diff --git a/examples/robotwin/robotwin_worker.py b/examples/robotwin/robotwin_worker.py
index ce1b66c..494458e 100644
--- a/examples/robotwin/robotwin_worker.py
+++ b/examples/robotwin/robotwin_worker.py
@@ -4,10 +4,11 @@
import argparse
import faulthandler
-import importlib
+import importlib.util
import json
import os
import signal
+import subprocess
import sys
import traceback
from pathlib import Path
@@ -15,6 +16,80 @@
import yaml
+def _normalize_pci_bus_id(value: str) -> str:
+ value = value.strip().lower()
+ domain, bus, device_function = value.split(":")
+ device, function = device_function.split(".", 1)
+ return (
+ f"{int(domain, 16):04x}:{int(bus, 16):02x}:"
+ f"{int(device, 16):02x}.{int(function, 16):x}"
+ )
+
+
+def _sapien_render_device_alias(gpu_id: str) -> str:
+ override = os.environ.get("STREAMINGWAM_SAPIEN_RENDER_DEVICE")
+ if override and override.lower() != "auto":
+ return override
+ completed = subprocess.run(
+ [
+ "nvidia-smi",
+ "-i",
+ str(gpu_id),
+ "--query-gpu=pci.bus_id",
+ "--format=csv,noheader",
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ return f"pci:{_normalize_pci_bus_id(completed.stdout.strip().splitlines()[0])}"
+
+
+def _bind_sapien3_renderer(gpu_id: str) -> None:
+ """Bind RoboTwin's Vulkan renderer to the requested physical GPU."""
+
+ import sapien as sapien_package
+ import sapien.core as sapien
+
+ version = str(getattr(sapien_package, "__version__", ""))
+ if not version.startswith("3."):
+ return
+ alias = _sapien_render_device_alias(gpu_id)
+ device = sapien.Device(alias)
+ if not device.can_render():
+ raise RuntimeError(f"SAPIEN device {alias} cannot render")
+ logged = False
+
+ def create_renderer(**_kwargs):
+ return sapien.render.SapienRenderer(device)
+
+ def create_scene(_engine, config=None):
+ nonlocal logged
+ config = sapien.SceneConfig() if config is None else config
+ sapien.physx.set_scene_config(config)
+ scene = sapien.Scene(
+ [sapien.physx.PhysxCpuSystem(), sapien.render.RenderSystem(device)]
+ )
+ selected = scene.render_system.device
+ if selected.pci_string != device.pci_string:
+ raise RuntimeError(
+ "SAPIEN render device mismatch: "
+ f"requested={device.pci_string}, selected={selected.pci_string}"
+ )
+ if not logged:
+ print(
+ f"[robotwin worker] renderer={selected.pci_string} "
+ f"cuda_id={selected.cuda_id}",
+ flush=True,
+ )
+ logged = True
+ return scene
+
+ sapien.SapienRenderer = create_renderer
+ sapien_package.SapienRenderer = create_renderer
+ sapien.Engine.create_scene = create_scene
+
+
def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser()
parser.add_argument("--gpu-id", required=True)
@@ -53,6 +128,37 @@ def write_job_status(path: Path, job: dict, *, phase: str, **extra) -> None:
_atomic_json(path, {**job, "phase": str(phase), **extra})
+def load_policy_module(policy_dir: Path):
+ """Load this checkout's deploy policy without top-level-name ambiguity."""
+
+ path = (policy_dir / "deploy_policy.py").resolve()
+ if not path.is_file():
+ raise FileNotFoundError(f"RoboTwin deploy policy does not exist: {path}")
+ module_name = f"_streamingwam_robotwin_deploy_policy_{os.getpid()}"
+ spec = importlib.util.spec_from_file_location(module_name, path)
+ if spec is None or spec.loader is None:
+ raise ImportError(f"cannot load RoboTwin deploy policy from {path}")
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ required = (
+ "get_model",
+ "eval",
+ "reset_model",
+ "prepare_instruction",
+ "prewarm_model",
+ "begin_timing_trajectory",
+ "end_timing_trajectory",
+ "_finish_episode",
+ )
+ missing = [name for name in required if not callable(getattr(module, name, None))]
+ if missing:
+ raise AttributeError(
+ f"deploy policy {path} is missing required hooks: {', '.join(missing)}"
+ )
+ print(f"[robotwin worker] deploy_policy={path}", flush=True)
+ return module
+
+
def main() -> int:
args = _parser().parse_args()
for stream in (sys.stdout, sys.stderr):
@@ -61,6 +167,7 @@ def main() -> int:
reconfigure(line_buffering=True)
faulthandler.register(signal.SIGUSR1, file=sys.stderr, all_threads=True)
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
+ _bind_sapien3_renderer(str(args.gpu_id))
robotwin = Path(args.robotwin_home).resolve()
policy_dir = Path(args.policy_dir).resolve()
output = Path(args.output).resolve()
@@ -72,9 +179,14 @@ def main() -> int:
str(policy_dir), str(robotwin), str(robotwin / "script"),
str(robotwin / "description" / "utils"),
]
- from eval_policy import class_decorator, eval_policy, get_camera_config, get_embodiment_config
+ import eval_policy as official_eval
+ from examples.robotwin.evaluator import evaluate_one_episode
+
+ class_decorator = official_eval.class_decorator
+ get_camera_config = official_eval.get_camera_config
+ get_embodiment_config = official_eval.get_embodiment_config
- policy_module = importlib.import_module("deploy_policy")
+ policy_module = load_policy_module(policy_dir)
model = policy_module.get_model({
"server_host": "127.0.0.1",
"server_port": args.server_port,
@@ -116,11 +228,16 @@ def environment_args(task: str, config: str) -> dict:
write_job_status(status_output, job, phase="environment_and_policy", gpu_id=args.gpu_id)
try:
task_environment = class_decorator(task)
- _, successes = eval_policy(
- task, task_environment, environment_args(task, config), model,
- 100000 * (1 + args.seed) + trial,
- test_num=1, video_size=None, instruction_type="unseen",
- skip_get_obs_within_replan=False,
+ _, successes = evaluate_one_episode(
+ official=official_eval,
+ task_name=task,
+ task_env=task_environment,
+ settings=environment_args(task, config),
+ model=model,
+ start_seed=100000 * (1 + args.seed) + trial,
+ trial=trial,
+ instruction_type="unseen",
+ policy_module=policy_module,
)
result = {
**job, "status": "completed", "success": int(successes),
@@ -138,6 +255,10 @@ def environment_args(task: str, config: str) -> dict:
policy_module._finish_episode(model)
model.close()
timing = list(model._timing_records) + list(model._episode_records)
+ for record in timing:
+ record.setdefault("task", task)
+ record.setdefault("config", config)
+ record.setdefault("trial", trial)
_atomic_json(output, {"gpu_id": args.gpu_id, "result": result, "timing": timing})
write_job_status(
status_output, job, phase="complete", gpu_id=args.gpu_id,
diff --git a/examples/robotwin/scripts/run_streamingwam_robotwin_benchmark.sh b/examples/robotwin/scripts/run_streamingwam_robotwin_benchmark.sh
new file mode 100644
index 0000000..6d4b3be
--- /dev/null
+++ b/examples/robotwin/scripts/run_streamingwam_robotwin_benchmark.sh
@@ -0,0 +1,84 @@
+#!/usr/bin/env bash
+# Canonical RoboTwin benchmark: fixed trials, R16, per-task success/Chunk/Total Time.
+# CHECKPOINT_FORMAT and MODE select released family defaults; explicit
+# NUM_INFERENCE_STEPS/ACTION_NUM_INFERENCE_STEPS always take precedence.
+set -euo pipefail
+
+ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)
+cd "$ROOT"
+
+MODE=${MODE:?set MODE=baseline|cd|ac-stream}
+CKPT=${CKPT:?set CKPT}
+CHECKPOINT_FORMAT=${CHECKPOINT_FORMAT:?set CHECKPOINT_FORMAT=fastwam|starwam|streamingwam}
+CONFIG=${CONFIG:?set CONFIG to the RoboTwin recipe}
+STATS_PATH=${STATS_PATH:?set STATS_PATH}
+BACKBONE_PATH=${BACKBONE_PATH:?set BACKBONE_PATH}
+ROBOTWIN_HOME=${ROBOTWIN_HOME:?set ROBOTWIN_HOME}
+INFERENCE_PYTHON=${INFERENCE_PYTHON:?set INFERENCE_PYTHON}
+SIMULATOR_PYTHON=${SIMULATOR_PYTHON:?set SIMULATOR_PYTHON}
+
+GPU_TOPOLOGY=${GPU_TOPOLOGY:-colocated}
+GPU_IDS=${GPU_IDS:-0,1,2,3}
+NUM_TRIALS=${NUM_TRIALS:-100}
+REPLAN_STEPS=${REPLAN_STEPS:-16}
+MODEL_SEED=${MODEL_SEED:-${SEED:-42}}
+EPISODE_SEED=${EPISODE_SEED:-$MODEL_SEED}
+TEXT_CACHE_PATH=${TEXT_CACHE_PATH:-/inspire/qb-ilm/project/qproject-fundationmodel/yangyi-253108120173/yzy/starwam/cache/text_embeds_cache}
+OUTPUT_DIR=${OUTPUT_DIR:-$ROOT/outputs/robotwin_${MODE}_${CHECKPOINT_FORMAT}_r${REPLAN_STEPS}_${NUM_TRIALS}trials_$(date +%Y%m%d_%H%M%S)}
+
+command=(
+ "$INFERENCE_PYTHON" examples/robotwin/multigpu_rollout.py
+ --config "$CONFIG"
+ --checkpoint "$CKPT"
+ --checkpoint-format "$CHECKPOINT_FORMAT"
+ --stats-path "$STATS_PATH"
+ --backbone-path "$BACKBONE_PATH"
+ --text-cache-path "$TEXT_CACHE_PATH"
+ --robotwin-home "$ROBOTWIN_HOME"
+ --inference-python "$INFERENCE_PYTHON"
+ --simulator-python "$SIMULATOR_PYTHON"
+ --inference-mode "$MODE"
+ --num-trials "$NUM_TRIALS"
+ --replan-steps "$REPLAN_STEPS"
+ --seed "$MODEL_SEED"
+ --episode-seed "$EPISODE_SEED"
+ --output-dir "$OUTPUT_DIR"
+)
+
+case "$GPU_TOPOLOGY" in
+ colocated)
+ # One inference server and one isolated simulator worker share each GPU.
+ # D8 action execution suppresses redundant SAPIEN renders while inference
+ # is in flight, so the four-GPU path keeps four jobs active in parallel.
+ command+=(--gpu-ids "$GPU_IDS")
+ ;;
+ split)
+ INFERENCE_GPU_IDS=${INFERENCE_GPU_IDS:-0,1}
+ SIMULATOR_GPU_IDS=${SIMULATOR_GPU_IDS:-2,3}
+ command+=(
+ --inference-gpu-ids "$INFERENCE_GPU_IDS"
+ --simulator-gpu-ids "$SIMULATOR_GPU_IDS"
+ --require-gpu-isolation
+ )
+ ;;
+ *)
+ echo "GPU_TOPOLOGY must be colocated or split, got: $GPU_TOPOLOGY" >&2
+ exit 2
+ ;;
+esac
+
+if [[ -n "${NUM_INFERENCE_STEPS:-}" ]]; then
+ command+=(--num-inference-steps "$NUM_INFERENCE_STEPS")
+fi
+if [[ -n "${ACTION_NUM_INFERENCE_STEPS:-}" ]]; then
+ command+=(--action-num-inference-steps "$ACTION_NUM_INFERENCE_STEPS")
+fi
+if [[ "$MODE" == "ac-stream" ]]; then
+ if [[ "${AC_STREAM_BACKEND:-accelerated}" == "eager" ]]; then
+ command+=(--ac-stream-eager)
+ else
+ command+=(--ac-stream-accelerated)
+ fi
+fi
+
+exec "${command[@]}" "$@"
diff --git a/examples/robotwin/timing.py b/examples/robotwin/timing.py
index 3e70f1f..1641ade 100644
--- a/examples/robotwin/timing.py
+++ b/examples/robotwin/timing.py
@@ -1,30 +1,160 @@
-"""Small, benchmark-facing RoboTwin timing aggregation helpers."""
+"""Benchmark-facing RoboTwin success, Chunk Time, and Total Time aggregation."""
from __future__ import annotations
+from collections import defaultdict
import statistics
from typing import Iterable
-def aggregate_evaluation(records: Iterable[dict]) -> dict:
- records = list(records)
- chunks = [
- float(record["model_inference_ms"])
+def _mean(values: list[float]) -> float | None:
+ return statistics.fmean(values) if values else None
+
+
+def _chunk_records(records: list[dict]) -> list[dict]:
+ inference = [
+ record
for record in records
if record.get("record_type") == "inference"
and not bool(record.get("warmup"))
and "model_inference_ms" in record
]
+ d8 = [record for record in inference if record.get("regime") == "d8"]
+ return d8 or inference
+
+
+def aggregate_evaluation(records: Iterable[dict]) -> dict:
+ """Aggregate fixed-episode results without hiding task-level variation.
+
+ Chunk Time is the CUDA-synchronized model call (D8 for AC-Stream, sync for
+ baseline/CD). Total Time uses successful trajectories only, averages each
+ task/config setting first, and then weights every covered setting equally.
+ Failed trajectories remain in the success-rate denominator.
+ """
+
+ records = list(records)
+ selected_chunks = _chunk_records(records)
episodes = [
record
for record in records
if record.get("record_type") == "episode" and "total_time_s" in record
]
- episode_times = [float(record["total_time_s"]) for record in episodes]
+ successful = [record for record in episodes if record.get("success") is True]
+
+ setting_episodes: dict[tuple[str, str], list[dict]] = defaultdict(list)
+ successful_times: dict[tuple[str, str], list[float]] = defaultdict(list)
+ setting_chunks: dict[tuple[str, str], list[float]] = defaultdict(list)
+ for record in episodes:
+ key = (str(record.get("config")), str(record.get("task")))
+ setting_episodes[key].append(record)
+ for record in successful:
+ key = (str(record.get("config")), str(record.get("task")))
+ successful_times[key].append(float(record["total_time_s"]))
+ for record in selected_chunks:
+ key = (str(record.get("config")), str(record.get("task")))
+ setting_chunks[key].append(float(record["model_inference_ms"]))
+
+ by_setting = {}
+ for key, selected in sorted(setting_episodes.items()):
+ config, task = key
+ successes = sum(record.get("success") is True for record in selected)
+ attempts = len(selected)
+ times = successful_times.get(key, [])
+ by_setting[f"{config}/{task}"] = {
+ "task": task,
+ "config": config,
+ "successes": successes,
+ "episodes": attempts,
+ "success_rate": successes / attempts if attempts else None,
+ "total_time_success_s": _mean(times),
+ "successful_timed_episodes": len(times),
+ "chunk_time_ms": _mean(setting_chunks.get(key, [])),
+ }
+
+ task_settings: dict[str, list[tuple[str, str]]] = defaultdict(list)
+ for key in setting_episodes:
+ task_settings[key[1]].append(key)
+ by_task = {}
+ for task, keys in sorted(task_settings.items()):
+ selected = [record for key in keys for record in setting_episodes[key]]
+ successes = sum(record.get("success") is True for record in selected)
+ attempts = len(selected)
+ setting_means = [
+ mean
+ for key in keys
+ if (mean := _mean(successful_times.get(key, []))) is not None
+ ]
+ by_task[task] = {
+ "successes": successes,
+ "episodes": attempts,
+ "success_rate": successes / attempts if attempts else None,
+ "total_time_success_s": _mean(setting_means),
+ "covered_configs": len(setting_means),
+ }
+
+ config_settings: dict[str, list[tuple[str, str]]] = defaultdict(list)
+ for key in setting_episodes:
+ config_settings[key[0]].append(key)
+ by_config = {}
+ for config, keys in sorted(config_settings.items()):
+ selected = [record for key in keys for record in setting_episodes[key]]
+ successes = sum(record.get("success") is True for record in selected)
+ attempts = len(selected)
+ setting_time_means = [
+ mean
+ for key in keys
+ if (mean := _mean(successful_times.get(key, []))) is not None
+ ]
+ config_chunk_values = [
+ value for key in keys for value in setting_chunks.get(key, [])
+ ]
+ by_config[config] = {
+ "successes": successes,
+ "episodes": attempts,
+ "success_rate": successes / attempts if attempts else None,
+ "total_time_success_s": _mean(setting_time_means),
+ "covered_settings": len(setting_time_means),
+ "chunk_time_ms": _mean(config_chunk_values),
+ }
+
+ setting_means = [
+ mean
+ for values in successful_times.values()
+ if (mean := _mean(values)) is not None
+ ]
+ chunks = [float(record["model_inference_ms"]) for record in selected_chunks]
+ all_inference = [
+ record
+ for record in records
+ if record.get("record_type") == "inference"
+ and not bool(record.get("warmup"))
+ and "model_inference_ms" in record
+ ]
+ by_regime = {
+ regime: _mean(
+ [
+ float(record["model_inference_ms"])
+ for record in all_inference
+ if record.get("regime") == regime
+ ]
+ )
+ for regime in ("sync", "d0", "d8")
+ if any(record.get("regime") == regime for record in all_inference)
+ }
+
return {
- "chunk_time_ms": statistics.fmean(chunks) if chunks else None,
- "total_time_per_episode_s": statistics.fmean(episode_times) if episode_times else None,
+ "chunk_time_ms": _mean(chunks),
+ "chunk_time_by_regime_ms": by_regime,
+ "total_time_per_episode_s": _mean(setting_means),
+ "successful_episode_time_s": _mean(
+ [float(record["total_time_s"]) for record in successful]
+ ),
"chunks": len(chunks),
- "successes": sum(record.get("success") is True for record in episodes),
+ "successes": len(successful),
"episodes": len(episodes),
+ "success_rate": len(successful) / len(episodes) if episodes else None,
+ "timed_settings": len(setting_means),
+ "by_config": by_config,
+ "by_task": by_task,
+ "by_setting": by_setting,
}
diff --git a/streamingwam/inference/ac_stream.py b/streamingwam/inference/ac_stream.py
index 52b7f72..4658e8e 100644
--- a/streamingwam/inference/ac_stream.py
+++ b/streamingwam/inference/ac_stream.py
@@ -538,6 +538,22 @@ def __init__(
def cursor(self) -> int:
return self._cursor
+ @property
+ def observation_required(self) -> bool:
+ """Whether the next action request launches an observation-conditioned chunk."""
+
+ if self._current is None:
+ return True
+ return (
+ self._future is None
+ and self._cursor >= self._next_launch_cursor
+ and self._cursor < self._window_end_cursor
+ )
+
+ @property
+ def inference_in_flight(self) -> bool:
+ return self._future is not None
+
def start_episode(self, observation: Any) -> ACStreamPrediction:
if self._closed:
raise RuntimeError("AC-Stream controller is closed")
diff --git a/tests/test_robotwin_inference_modes.py b/tests/test_robotwin_inference_modes.py
index a2c592e..a5b874f 100644
--- a/tests/test_robotwin_inference_modes.py
+++ b/tests/test_robotwin_inference_modes.py
@@ -5,7 +5,10 @@
import torch
from streamingwam.eval.policy import StreamingWAMPolicy
+from streamingwam.inference.ac_stream import ACStreamController, ACStreamPrediction
from examples.robotwin.runtime import resolve_inference_runtime
+from examples.robotwin.multigpu_rollout import resolve_sampling_steps
+from examples.robotwin.timing import aggregate_evaluation
class _RecordingModel:
@@ -109,3 +112,73 @@ def test_robotwin_runtime_rejects_ac_stream_backend_for_sync_modes(mode: str) ->
def test_robotwin_runtime_rejects_conflicting_ac_stream_backends() -> None:
with pytest.raises(ValueError, match="mutually exclusive"):
resolve_inference_runtime("ac-stream", accelerated=True, eager=True)
+
+
+@pytest.mark.parametrize(
+ ("checkpoint_format", "mode", "expected"),
+ [
+ ("starwam", "baseline", (4, 4)),
+ ("starwam", "cd", (1, 1)),
+ ("starwam", "ac-stream", (1, 1)),
+ ("fastwam", "baseline", (10, 10)),
+ ("fastwam", "cd", (1, 2)),
+ ("streamingwam", "ac-stream", (1, 2)),
+ ],
+)
+def test_robotwin_sampling_steps_follow_checkpoint_family(
+ checkpoint_format: str,
+ mode: str,
+ expected: tuple[int, int],
+) -> None:
+ assert resolve_sampling_steps(checkpoint_format, mode, None, None) == expected
+
+
+def test_robotwin_sampling_step_overrides_are_preserved() -> None:
+ assert resolve_sampling_steps("fastwam", "baseline", 3, 5) == (3, 5)
+
+
+def test_ac_stream_requests_observation_only_at_d0_and_d8_launch() -> None:
+ calls = []
+
+ def predict(observation, previous, delay):
+ calls.append((observation, previous, delay))
+ actions = torch.zeros((32, 7), dtype=torch.float32)
+ return ACStreamPrediction(
+ env_actions=actions.numpy(),
+ model_actions=actions,
+ communication_ms=0.0,
+ inference_ms=1.0,
+ )
+
+ controller = ACStreamController(predict)
+ controller.start_episode("d0")
+ assert controller.observation_required is False
+ for _ in range(8):
+ controller.next_action(None)
+ controller.mark_action_executed()
+ assert controller.observation_required is True
+ controller.next_action("d8")
+ controller.mark_action_executed()
+ assert controller.observation_required is False
+ controller.close()
+ assert [call[2] for call in calls] == [0, 8]
+
+
+def test_robotwin_summary_uses_success_only_macro_total_and_d8_chunk() -> None:
+ records = [
+ {"record_type": "inference", "regime": "d0", "model_inference_ms": 100.0},
+ {"record_type": "inference", "regime": "d8", "model_inference_ms": 40.0},
+ {"record_type": "inference", "regime": "d8", "model_inference_ms": 50.0},
+ {"record_type": "episode", "task": "task", "config": "clean", "success": True, "total_time_s": 10.0},
+ {"record_type": "episode", "task": "task", "config": "clean", "success": True, "total_time_s": 30.0},
+ {"record_type": "episode", "task": "task", "config": "clean", "success": False, "total_time_s": 999.0},
+ {"record_type": "episode", "task": "task", "config": "random", "success": True, "total_time_s": 100.0},
+ ]
+
+ summary = aggregate_evaluation(records)
+
+ assert summary["chunk_time_ms"] == pytest.approx(45.0)
+ assert summary["total_time_per_episode_s"] == pytest.approx(60.0)
+ assert summary["successes"] == 3
+ assert summary["episodes"] == 4
+ assert summary["by_setting"]["clean/task"]["total_time_success_s"] == pytest.approx(20.0)