Add Cosmos3 Policy#894
Conversation
Signed-off-by: Riddler <luya.wang@qq.com>
Greptile SummaryThis PR adds a Cosmos3 remote policy extension. The main changes are:
Confidence Score: 4/5The Cosmos3 policy path is not safe to merge until the setup, construction, and cleanup contracts are fixed.
isaaclab_arena_cosmos3/docker/COSMOS3_COMMIT; isaaclab_arena_cosmos3/policy/cosmos3_remote_policy.py Important Files Changed
Reviews (1): Last reviewed commit: "Add openpi-client runtime dependency and..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -0,0 +1 @@ | |||
| TODO: pin a known-working cosmos-framework commit | |||
There was a problem hiding this comment.
The default server script reads this file and runs git checkout "$PINNED_COMMIT". With the placeholder text here, the documented ./run_cosmos3_server.sh path clones cosmos-framework and then fails before building the inference server.
| name = "cosmos3_remote" | ||
| config_class = Cosmos3RemotePolicyArgs | ||
|
|
||
| def __init__(self, config: Cosmos3RemotePolicyArgs, cosmos3_embodiment_adapter: Cosmos3EmbodimentAdapter) -> None: |
There was a problem hiding this comment.
Registered policies are constructed from their config as policy_type(policy_cfg). This constructor requires a second cosmos3_embodiment_adapter argument, so selecting cosmos3_remote through the typed runner path raises TypeError before rollout starts.
| @property | ||
| def is_remote(self) -> bool: | ||
| return True |
There was a problem hiding this comment.
Returning True marks this policy for remote-policy teardown, where the runner calls shutdown_remote() on remote policies. This class only implements close(), so a rollout can finish and then fail during cleanup with an AttributeError.
| assert ( | ||
| chunk.ndim == 2 and chunk.shape[1] == self._cosmos3_embodiment_adapter.action_dim | ||
| ), f"Expected action of shape (H, {self._cosmos3_embodiment_adapter.action_dim}); got {chunk.shape}" | ||
| assert ( | ||
| chunk.shape[0] >= self._open_loop_horizon | ||
| ), f"Server returned horizon {chunk.shape[0]} < configured open_loop_horizon {self._open_loop_horizon}" |
There was a problem hiding this comment.
These checks validate data from the WebSocket server, but assert statements are removed under python -O. In that mode, a wrong-width or too-short action response can enter the cache and later produce malformed actions or indexing failures instead of a clear policy error.
| assert self.task_description, ( | ||
| "Cosmos3RemotePolicy requires a non-empty language instruction" | ||
| " (set via --language_instruction or on the task definition)." | ||
| ) |
There was a problem hiding this comment.
This language-instruction guard also disappears under python -O. If the environment returns an empty or missing task description, the policy sends that value in prompt to the remote server instead of stopping locally with a clear error.
| right = right_t.squeeze(0).permute(1, 2, 0).numpy().astype(wrist.dtype) | ||
|
|
||
| # Compose: vstack(wrist, hstack(left, right)) → (720, 640, 3) | ||
| image = np.concatenate((wrist, np.concatenate((left, right), axis=1)), axis=0) |
There was a problem hiding this comment.
This sends the wrist view on the top half and the two exterior views below it. If the Cosmos3 DROID server was trained on a different packed layout, every rollout still produces a valid (720, 640, 3) image but the model reads each camera from the wrong spatial position and returns bad actions without an obvious error.
![arena-review-bot[bot]](https://avatars.githubusercontent.com/u/195794629?s=60&v=4){kind=small}
| wrist = image_tools.resize_with_pad(extracted.wrist_image, self.image_h, self.image_w) | ||
|
|
||
| # Resize left/right to full size first, then bilinear downsample. | ||
| half_size = (self.image_h // 2, self.image_w // 2) # (180, 320) |
There was a problem hiding this comment.
🔴 The composed image is 540×640, not the 720×640 this method promises. resize_with_pad makes wrist (360, 640); half_size=(180, 320) makes left/right (180, 320) each, so hstack(left, right) is (180, 640) and the final vstack is (540, 640, 3). That contradicts the class docstring ("single 720×640 image"), this method's docstring, the assert message on line 97, and test_droid_adapter_uses_cosmos3_wire_keys, which asserts (720, 640, 3) — so that test fails as written. The assert expression on line 96 (image_h + image_h//2) was written to match the 540 result, so it passes and hides the mismatch. If cosmos3's wire format really is 720×640 (wrist on top, left|right filling the bottom half), the bottom row needs full height:
| half_size = (self.image_h // 2, self.image_w // 2) # (180, 320) | |
| half_size = (self.image_h, self.image_w // 2) # (360, 320) |
and line 96's assert should become (2 * self.image_h, self.image_w, 3). Could you confirm the server's expected layout and reconcile all four spots? A wrong image shape silently corrupts every cosmos3 eval.
|
|
||
|
|
||
| @register_policy | ||
| class Cosmos3RemotePolicy(PolicyBase): |
There was a problem hiding this comment.
🟡 This file is a near-verbatim clone of isaaclab_arena_openpi/policy/pi0_remote_policy.py — the __init__ cache setup, get_action, reset, _maybe_init_per_env_state, _call_server_with_retry, _close_websocket_best_effort, and the _resolve_*_adapter pattern are all mechanically renamed copies. The # TODO(cvolk...): add a RemotePolicy base class to unify this and other remote policies in from_dict is exactly the abstraction this duplication is asking for — and the same TODO already sits in pi0. Could we pull the chunk-replay cache + retry/reconnect machinery into a RemotePolicy base in isaaclab_arena/policy/ now, so cosmos3 and pi0 only supply the response key (action vs actions), the optional post-process, and adapter resolution? Otherwise the two paths drift silently. (Same theme, minor: Cosmos3RemotePolicyArgs breaks the ...Cfg / PolicyCfg-subclass convention that Pi0RemotePolicyCfg follows, and PolicyBase is left unparameterized here where pi0 uses PolicyBase[Pi0RemotePolicyCfg].)
| remote_port: int = 8000 | ||
| """Port the cosmos3 inference server listens on.""" | ||
|
|
||
| num_envs: int = 1 |
There was a problem hiding this comment.
🔵 num_envs is never read — the per-env caches are lazy-allocated from env.unwrapped.num_envs in _maybe_init_per_env_state, not from this field, so the docstring's "Used to pre-allocate per-environment action caches" is inaccurate too. Suggest dropping the field (it isn't wired through add_args_to_parser/from_args either) unless a later change needs it.
| @@ -0,0 +1 @@ | |||
| TODO: pin a known-working cosmos-framework commit | |||
There was a problem hiding this comment.
🟡 This placeholder is read literally by run_cosmos3_server.sh (PINNED_COMMIT=$(tr -d ... < COSMOS3_COMMIT) → git checkout "$PINNED_COMMIT"), so the default build path runs git checkout 'TODO: pin a known-working cosmos-framework commit' and fails. The server launcher can't work until a real commit is pinned — is this meant to land before merge, or is the PR waiting on it?
🤖 Isaac Lab-Arena Review BotSummaryThis PR adds a new Design, Boundaries & Scope
Findings🔴 Critical: droid_adapter.py:82 — Composed image is (540, 640, 3), not the 720×640 the class/method docstrings, the assert message, and the unit test all claim; the assert expression was written to match 540, so it silently passes. The wire-format image is malformed → corrupts every cosmos3 eval. Test CoverageGood client-side coverage (caching/advancement, parallel envs, reset-by-env-ids, gripper binarization, retry/reconnect, shape asserts). These are non-sim CPU tests, so the inner/outer VerdictMinor fixes needed — resolve the image-shape bug (blocking) and ideally the duplication before merge. |
Correct composed image shape from (720, 640, 3) to (540, 640, 3) in class and method docstrings, assert message, and unit test to match the actual wire format used by cosmos3 DROID adapter. Pin COSMOS3_COMMIT to a known-working Release (2026-07-11) so the server Docker build succeeds without user intervention. Delete unused num_envs field from Cosmos3RemotePolicyArgs. Signed-off-by: ptrdlr14 <ptrdlr14@gmail.com>
Consistent with pi0 and dreamzero: use config.policy_device instead of env.unwrapped.device, and expose --policy_device CLI argument. Signed-off-by: ptrdlr14 <ptrdlr14@gmail.com>
1 participant
Summary
Add Cosmos3 Policy support via
openpi-server.Detailed description
Introduced isaaclab_arena_cosmos3 Extension: Created a new extension package dedicated to Cosmos3 policy integration.
Added OpenPI Client & Policy: Implemented Cosmos3RemotePolicy which establishes a WebSocket connection to a remote inference server using
openpi_client.