Add isaaclab_arena_openpi package for pi0 droid eval

[cvolkcvolk]

Summary

Adds isaaclab_arena_openpi/, a new top-level package that lets arena evaluate the upstream openpi pi0/pi05/pi0_fast droid policies. For convenience ships the current openpi server as a dockerized container fixed to current main.

What's in the package

• policy/pi0_droid_remote_policy.py — Pi0DroidRemotePolicy (PolicyBase subclass). Connects to an openpi server, fetches one action chunk per open_loop_horizon steps, reconnects on dropped websockets up to MAX_RECONNECT_ATTEMPTS. --Single-env only (openpi server is single-obs, asserted at runtime).

• policy/pi0_droid_config.py — wire-format constants, the variant horizon table, and Pi0DroidRemotePolicyArgs.

• docker/build_openpi_server.sh — clones upstream Physical-Intelligence/openpi at a pinned commit, builds the server image with upstream's own serve_policy.Dockerfile. Optionally pushes to ngc for use on osmo/ nightly etc.

Usage

./isaaclab_arena_openpi/docker/build_openpi_server.sh    # one-time, ~3 min, ~19 GB

# Terminal 1: start server (pi05 default)
docker run --rm -it --gpus all --network=host \
  -e XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 \
  isaaclab_arena_openpi-server:latest \
  uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_droid_jointpos_polaris \
    --policy.dir=gs://openpi-assets-simeval/pi05_droid_jointpos

# Terminal 2: arena client (inside the arena docker container)
python isaaclab_arena/evaluation/policy_runner.py     --viz kit     --policy_type isaaclab_arena_openpi.policy.pi0_droid_remote_policy.Pi0DroidRemotePolicy     --policy_variant pi05     --remote_host localhost --remote_port 8000     --num_steps 500     --enable_cameras --num_envs 1     --language_instruction "Pick up the Rubik's cube and place it in the bowl."     pick_and_place_maple_table       --embodiment droid_abs_joint_pos       --pick_up_object rubiks_cube_hot3d_robolab       --destination_location bowl_ycb_robolab       --hdr home_office_robolab

[isaaclab-review-bot]

Review: isaaclab_arena_openpi Package

Thanks for adding this openpi integration! The overall architecture is clean and follows the established patterns from isaaclab_arena_gr00t. I've reviewed the implementation and have a few observations:

✅ Strengths

1. Clean separation of concerns - Wire-format constants and variant table in pi0_droid_config.py, execution logic in the policy module. This makes maintenance straightforward when upstream changes.

2. Well-documented - The README thoroughly explains the layout, installation, variant table, and rationale for not bundling an environment. The inline docstrings in the policy are also helpful.

3. Comprehensive unit tests - 13 tests covering wire format, chunk caching, binarisation, clipping, reconnect logic, and edge cases like short server responses. The stubbing approach is clever and allows testing without heavy dependencies.

4. Defensive chunking - The cache-and-advance pattern with explicit flush on reconnect is a sensible approach for a single-env scenario.

---

🔍 Observations / Questions

1. Broad exception handling in reconnect logic

_call_server_with_retry catches OSError in addition to the websocket-specific exceptions:

retryable_exceptions = (
    websockets.exceptions.ConnectionClosedError,
    websockets.exceptions.ConnectionClosedOK,
    OSError,
)

OSError is quite broad and could potentially mask unrelated errors (e.g., file permission issues if any file I/O is involved). Consider narrowing to specific socket errors like ConnectionRefusedError and ConnectionResetError, or at least logging which specific exception triggered the retry.

2. Missing validation for numeric edge cases

_unpack_pi0_response validates shape but doesn't check for NaN/Inf values which could propagate silently to the action space:

def _unpack_pi0_response(self, server_response: dict[str, Any]) -> np.ndarray:
    action_chunk = np.asarray(server_response[WIRE_ACTIONS])
    # ... shape validation only

A quick np.isfinite check could catch malformed server responses early.

3. Test coverage gap: empty task description

set_task_description raises ValueError for empty/None input, but this path isn't covered by tests:

def set_task_description(self, task_description: str | None) -> str:
    if not task_description:
        raise ValueError(...)

Consider adding a quick test for this.

4. Minor: Hardcoded gripper binarisation threshold

The 0.5 threshold is noted as matching upstream checkpoint behavior, but if a future checkpoint changes this, it would require a code change. Could be added to Pi0Variant for extensibility, though this may be over-engineering for now.

---

📝 Note on PR Description

The PR description mentions changes to PolicyBase.from_args and PolicyBase.add_args_to_parser with derived defaults, but those changes don't appear in this diff. Is that intentional for a follow-up, or did they get dropped?

---

Overall this is solid work for a draft. Once the end-to-end tests against a live openpi server are validated, this should be good to go. 👍

[isaaclab-review-bot]

Updated Review (Post-Cleanup #2)

Thanks for addressing the S3 concern! The latest commits simplify documentation and fix the pi0_fast bucket issue.

✅ Changes Since Last Review (31ee8a4 → 265b989)

1. Fixed pi0_fast bucket - Changed from s3://openpi-assets-simeval/... to gs://openpi-assets/checkpoints/polaris/pi0_fast_droid_jointpos_polaris. Now all variants use GCS (gs://), no mixed storage backends! 🎉

2. README simplification - Removed verbose explanation about checkpoint selection rationale (cleaner docs)

3. Config/policy comment cleanup - Trimmed module docstrings and inline comments in pi0_droid_config.py and pi0_droid_remote_policy.py (less noise, same functionality)

4. Minor code style - Changed raise ValueError → assert for variant validation

---

⚠️ One Small Note

The raise ValueError → assert change in pi0_droid_remote_policy.py means the check can be optimized away with python -O. If this is intentional (trusting callers), fine. If the validation should always run, consider keeping ValueError.

---

✅ Previous Review Highlights (Still Valid)

• Clean action pipeline with gripper binarisation delegated to arena's config
• Docker server build workflow for reproducibility
• Polaris config variants with proper GCS checkpoints

---

✅ Overall

All previous concerns addressed. Cleanup looks good. Ready to merge! 👍

---

📌 Reviewed at: 265b989

[cvolkcvolk]

For now I put the openpi dependency on the client into the arena base docker. We could in the future layer this instead ontop as a new container

[greptile-apps]

Greptile Summary

Adds isaaclab_arena_openpi/, a new top-level package that wraps the upstream openpi WebSocket server as a PolicyBase subclass (Pi0DroidRemotePolicy) for evaluating pi0/pi05/pi0_fast DROID policies inside Isaac Lab Arena, along with a Docker build script that pins the server to a matching commit and unit tests covering the wire-format contract and reconnect logic.

• Pi0DroidRemotePolicy fetches action chunks over WebSocket, caches them for open_loop_horizon steps, and reconnects with retry on dropped connections; single-env only, enforced by assertion.
• build_openpi_server.sh clones upstream openpi at a pinned SHA, appends a COPY . /app layer to the upstream Dockerfile, and optionally pushes to NGC — the pinned commit matches OPENPI_COMMIT in the arena Dockerfile so client and server share the same wire format.
• Dockerfile.isaaclab_arena installs openpi-client at the same pinned commit via --no-deps pip install, consistent with other sub-package additions.

Confidence Score: 4/5

Safe to merge; the new package is self-contained and does not modify existing evaluation paths.

The implementation is clean and well-tested. The only notable items are a docstring that overstates the reconnect count by one and no explicit close/cleanup path for the WebSocket client resource. Neither affects correctness during a normal eval run.

pi0_droid_remote_policy.py — the retry docstring and missing WebSocket cleanup are worth a quick look before this pattern is copied to future remote policies.

Important Files Changed

Filename	Overview
isaaclab_arena_openpi/policy/pi0_droid_remote_policy.py	New Pi0DroidRemotePolicy class with WebSocket-based action fetching, open-loop chunk caching, and retry logic; docstring on _call_server_with_retry overstates reconnect count by one, and no close() path exists for the WebSocket resource.
isaaclab_arena_openpi/policy/pi0_droid_config.py	Clean constants and dataclass config file; variant horizon table is straightforward and easy to extend.
isaaclab_arena_openpi/tests/test_pi0_droid_remote_policy.py	Well-structured unit tests covering wire-format keys, chunk caching/indexing, reconnect recovery, and give-up-after-max-attempts; patching strategy correctly targets _wait_for_server to avoid real connections.
isaaclab_arena_openpi/docker/build_openpi_server.sh	Correct pinned-commit Docker build workflow; does a full git clone that could be made faster with --filter=blob:none.
docker/Dockerfile.isaaclab_arena	Adds pinned openpi-client installation and COPY for the new package directory; consistent with how other sub-packages are included.
setup.py	Trivial addition of isaaclab_arena_openpi* to the find_packages list; matches the pattern used by other sub-packages.

Sequence Diagram

sequenceDiagram
    participant Runner as policy_runner.py
    participant Policy as Pi0DroidRemotePolicy
    participant Client as WebsocketClientPolicy
    participant Server as openpi server (Docker)

    Runner->>Policy: set_task_description(instruction)
    loop Every open_loop_horizon steps
        Runner->>Policy: get_action(env, observation)
        Policy->>Policy: extract DROID obs (cameras + proprio)
        Policy->>Policy: pack pi0 wire-format request
        Policy->>Client: infer(request)
        Client-->>Server: WebSocket send
        Server-->>Client: action chunk (H x 8)
        Client-->>Policy: actions ndarray
        Policy->>Policy: cache chunk[:open_loop_horizon]
        Policy-->>Runner: action[0] as torch.Tensor (1x8)
    end
    loop Cached steps (open_loop_horizon - 1 remaining)
        Runner->>Policy: get_action(env, observation)
        Policy-->>Runner: cached action[i] as torch.Tensor (1x8)
    end
    note over Client,Server: On ConnectionClosedError/OSError: reconnect up to MAX_RECONNECT_ATTEMPTS-1 times, flush cache on each reconnect

Loading

_{Reviews (1): Last reviewed commit: "Address review feedback: trim docstrings..." | Re-trigger Greptile}

[greptile-apps]

The docstring claims to reconnect "up to MAX_RECONNECT_ATTEMPTS times" but the loop structure only reconnects MAX_RECONNECT_ATTEMPTS - 1 (2) times — the loop has 3 iterations, but the last iteration raises immediately on failure without reconnecting. A caller or operator who reads the docstring and tunes MAX_RECONNECT_ATTEMPTS will get one fewer reconnect than expected.

Suggested change

	def _call_server_with_retry(self, server_request: dict[str, Any]) -> dict[str, Any]:
	"""Send the request, reconnecting up to ``MAX_RECONNECT_ATTEMPTS`` times.
	def _call_server_with_retry(self, server_request: dict[str, Any]) -> dict[str, Any]:
	"""Send the request, making up to ``MAX_RECONNECT_ATTEMPTS`` total attempts
	(i.e. reconnecting up to ``MAX_RECONNECT_ATTEMPTS - 1`` times).

[greptile-apps]

Missing WebSocket cleanup on policy teardown

_websocket_client is created in __init__ but there is no close() or __del__ override to shut the connection down. If the runner exits or a new policy is instantiated in the same process, the underlying TCP connection and async event-loop resources held by WebsocketClientPolicy are leaked. Consider adding a close() (or __del__) that calls the equivalent cleanup method the upstream client exposes so the resource is freed when the policy is discarded.

[greptile-apps]

Full-history clone is slower than necessary

git clone without --filter=blob:none fetches the complete object graph including every historical blob. Consider adding --filter=blob:none (partial clone — fetches tree and commit objects only, blobs fetched on demand) to shorten the one-time build step without changing correctness.

[xyao-nv]

Can you unify this spec defining the src of openpi with Dockerfile? I wonder if having a shared .env could work.

[xyao-nv]

Can it be doc rst instead of readme?

[xyao-nv]

External users may not have write access to this path

[xyao-nv]

I wonder if the clone could be optional as users may already have a PI src repo and just want to build server docker (if Pi did not ship it).

[xyao-nv]

If we imagine every policy builds their own docker, then can we make the docker build & push more universal?
Not bind to Pi server docker.

I personally think it's a valid assumption.

[xyao-nv]

Can you add a TODO for env_id reset when parallel is supported?

[xyao-nv]

Can you add a TODO for extending to parallel envs?

[xyao-nv]

Does it have to droid dependent? Can you extend to other embodiments?
Can the dict key be a dataclass/cfg?

[xyao-nv]

We kinda enforce "camera_obs" as the key when inserting camera to embodiment. But how about "policy"? Is it fixed?

[xyao-nv]

I don't like the idea of limiting the class to only robot specific closedloop policy.
It's okay that you have only tested with Droid, but when building the policy class, it shall be embodiment agnostic, i.e. Pi0RemotePolicy.
You can provide embodiment specific cfg as the data member.

[xyao-nv]

Sorry I have to request changes -- can policy class be embodiment agostic?

[xyao-nv]

When closing the connection (esp it's used in eval_runner), what's the expected components needed? Gr00t side I have added closing socket and cleaning context to prevent stalling memory.

[xyao-nv]

I wonder if this func could be in the parent class RemotePolicy. Retry with max attempts will always be needed.

[xyao-nv]

Do you need action_scheduler_cls? This MR seems to only contain obs -> action_chunk. Can you add a TODO for action_chunk to action

[xyao-nv]

Shall we do a closedloop-embodiment dataclass? Hydra?
So we can unify how cfg are supposed to set, you have py, and g00t has yaml.