Server–Client Policies

Robot foundation models such as vision-language-action models (VLAs) and world-model policies bring their own dependencies, GPUs, and release schedules. Arena keeps them out of the simulation process. The model runs as a server, and Arena runs a thin client policy that implements PolicyBase. The two talk over a websocket, so the server can run in another container, on another GPU, or on another machine.

┌──────────── Arena process ────────────┐            ┌──── model server ────┐
│ env obs ─► EmbodimentAdapter          │  request   │                      │
│            extract + pack_request ────┼──────────► │  infer(request)      │
│                                       │            │                      │
│ env.step ◄─ action ◄─ chunk replay ◄──┼─────────── │  {"actions": (H, D)} │
└───────────────────────────────────────┘  response  └──────────────────────┘

Arena includes clients for Cosmos (CosmosRemotePolicy), π0 and π0.5 from OpenPI (Pi0RemotePolicy), GR00T (Gr00tRemoteClosedloopPolicy), and DreamZero (DreamZeroRemotePolicy). The OpenPI and GR00T quickstarts start a server and run its client end to end.

How the client works

The Cosmos and OpenPI clients share one implementation, remote_policy_base.py, which separates what is sent from how it is sent.

An EmbodimentAdapter decides what is sent for one robot and one server format. Its extract(obs, env_id) method pulls one environment’s camera images and joint state out of Arena’s observation dictionary. pack_request(extracted, instruction) turns them into the request the server expects. action_dim declares the width of each action the server returns.

RemoteChunkReplayPolicy decides how it is sent, and every adapter shares it. For each environment, it sends one request and receives an action chunk of shape (H, action_dim). It replays the first open_loop_horizon actions, one per step, and then requests a new chunk. When an environment resets, or the connection drops and reconnects, the client discards the cached chunk so that it never replays stale actions. close() disconnects the client but leaves the server running.

Requests and responses are flat dictionaries of NumPy arrays and strings, packed with msgpack_numpy. For example, the Cosmos DROID adapter sends:

{
    "observation/image": uint8[H, W, 3],          # wrist view on top, two exterior views below
    "observation/joint_position": float32[7],
    "observation/gripper_position": float32[1],
    "prompt": "put the banana in the bowl",
}
# response: {"action": float32[H, 8]}

The client’s configuration holds the connection settings: remote_host, remote_port, ping_interval, ping_timeout, and policy_device. The RoboLab example below evaluates π0.5 and Cosmos on the same two tasks:

Configuration file (robolab_2_tasks_pi0_and_cosmos.yaml)

# Runs two robolab tasks with both the OpenPI (pi0) and Cosmos policies (4 runs)
# Values under shared are used by every Run below.
# A Run only needs to list what is different; its values override the shared ones.
# To change a shared value for every Run on the command line:
# - with experiment_runner.py, use shared.<path>=<value>
# - with osmo/submit_arena_experiment.py, use experiment_cfg.shared.<path>=<value>
shared:
  environment:
    enable_cameras: true
  environment_builder:
    num_envs: 25
  rollout_limit:
    num_episodes: 100
  variations:
    droid_abs_joint_pos:
      camera_extrinsics_wrist_camera:
        enabled: false
        sampler_cfg:
          low: [-0.05, -0.05, -0.05]
          high: [0.05, 0.05, 0.05]

runs:

  banana_in_bowl_pi0:
    environment: &banana_in_bowl_env
      type: isaaclab_arena_environments/robolab/tasks/banana_in_bowl.yaml
    policy: &openpi_policy
      type: isaaclab_arena_openpi.policy.pi0_remote_policy.Pi0RemotePolicy
      policy_variant: pi05
      policy_device: cuda:0
      remote_host: 127.0.0.1
      remote_port: 8000
      openpi_embodiment_adapter: droid

  banana_in_bowl_cosmos:
    environment: *banana_in_bowl_env
    policy: &cosmos_policy
      type: isaaclab_arena_cosmos.policy.cosmos_remote_policy.CosmosRemotePolicy
      cosmos_embodiment_adapter: droid
      policy_device: cuda:0
      remote_host: 127.0.0.1
      remote_port: 8000

  banana_on_plate_pi0:
    environment: &banana_on_plate_env
      type: isaaclab_arena_environments/robolab/tasks/banana_on_plate.yaml
    policy: *openpi_policy

  banana_on_plate_cosmos:
    environment: *banana_on_plate_env
    policy: *cosmos_policy

GR00T and DreamZero servers use their own protocols, so Arena keeps their native clients. Both still implement PolicyBase, and the runners treat them like any other policy. GR00T also converts chunks into per-step actions through a configurable ActionScheduler (see action_scheduling/).

Add a model or a robot

To run an existing model on a new robot, write an EmbodimentAdapter for that robot’s cameras, joint state, and action layout. Then add its name to the client’s adapter lookup, which cosmos_embodiment_adapter or openpi_embodiment_adapter selects. The transport does not change.

To add a model whose server speaks the OpenPI protocol, subclass RemoteChunkReplayPolicy and PolicyBase[YourCfg]. Set name and server_response_actions_key, and pass your adapter and horizon to super().__init__. CosmosRemotePolicy is the reference: it differs from the OpenPI client only in its response key and adapter.

cosmos_remote_policy.py

from __future__ import annotations

from typing import TYPE_CHECKING

from isaaclab_arena.assets.register import register_policy
from isaaclab_arena.policy.policy_base import PolicyBase
from isaaclab_arena.policy.remote_policy_base import RemoteChunkReplayPolicy
from isaaclab_arena_cosmos.policy.cosmos_remote_config import CosmosRemotePolicyCfg

if TYPE_CHECKING:
    from isaaclab_arena_cosmos.policy.droid_adapter import CosmosDroidAdapter

# Cosmos serves over openpi's WebsocketPolicyServer, so this reuses the shared
# openpi-protocol machinery (RemoteChunkReplayPolicy). It differs only in the response key
# (the server returns the chunk under the singular "action") and the DROID wire format
# (see CosmosDroidAdapter); the open-loop horizon comes from the config rather than a
# per-variant table because the released Cosmos policies are a single DROID family.


@register_policy
class CosmosRemotePolicy(RemoteChunkReplayPolicy, PolicyBase[CosmosRemotePolicyCfg]):
    """Cosmos remote closed-loop policy, parameterized by an embodiment adapter."""

    name = "cosmos_remote"
    server_response_actions_key = "action"

    def __init__(self, config: CosmosRemotePolicyCfg) -> None:
        adapter = _resolve_cosmos_embodiment_adapter(config.cosmos_embodiment_adapter)
        super().__init__(config, adapter, config.open_loop_horizon)

    @classmethod
    def default_image_size(cls, config: CosmosRemotePolicyCfg) -> tuple[int, int]:
        """Return the (height, width) the selected Cosmos adapter resizes camera images to."""
        adapter = _resolve_cosmos_embodiment_adapter(config.cosmos_embodiment_adapter)
        return adapter.IMAGE_H, adapter.IMAGE_W


def _resolve_cosmos_embodiment_adapter(key: str) -> CosmosDroidAdapter:
    """Instantiate the adapter registered under ``key``.

    Imports are deferred to call time so adapter modules can import the shared
    ``EmbodimentAdapter`` at module top without creating a circular import.
    """
    if key == "droid":
        from isaaclab_arena_cosmos.policy.droid_adapter import CosmosDroidAdapter

        return CosmosDroidAdapter()
    raise ValueError(f"Unknown cosmos_embodiment_adapter {key!r}; expected 'droid'")

If the server uses its own protocol, implement PolicyBase directly, as the GR00T client does. Keep the client thin. Image preprocessing, chunk handling, and reconnects belong in the client, and model code belongs in the server.

Because the model sits behind a socket, scaling up is a deployment choice, not a code change. Point remote_host at a larger GPU or another node, and raise num_envs to run more environments per client. On OSMO, each client policy is paired with a server task in osmo/tasks/, so one workflow launches both. To serve a new model there, subclass PolicyServerTask and register it with @register_server_task.

Next steps

Continue to Agent as Policy to use a hosted foundation model as a policy without robot training. See Arena Experiments to submit a server–client evaluation to OSMO.