DRAFT Openpi client

isaaclab_arena/environments/arena_env_builder.py

@@ -289,12 +289,17 @@ def build_registered(
         )
         return name, cfg
 
-    def make_registered(self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None) -> ManagerBasedEnv:
-        env, _ = self.make_registered_and_return_cfg(env_cfg)
+    def make_registered(
+        self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None, render_mode: str | None = None
+    ) -> ManagerBasedEnv:
+        env, _ = self.make_registered_and_return_cfg(env_cfg, render_mode=render_mode)
         return env
 
     def make_registered_and_return_cfg(
-        self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None
+        self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None, render_mode: str | None = None
     ) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
         name, cfg = self.build_registered(env_cfg)
-        return gym.make(name, cfg=cfg).unwrapped, cfg
+        env = gym.make(name, cfg=cfg, render_mode=render_mode)
+        if render_mode is None:
+            env = env.unwrapped
+        return env, cfg

isaaclab_arena/evaluation/eval_runner.py

@@ -8,6 +8,7 @@
 import json
 import os
 import traceback
+from gymnasium.wrappers import RecordVideo
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
@@ -23,7 +24,7 @@
     from isaaclab_arena.policy.policy_base import PolicyBase
 
 
-def load_env(arena_env_args: list[str], job_name: str):
+def load_env(arena_env_args: list[str], job_name: str, render_mode: str | None = None):
 
     reload_arena_modules()
 
@@ -38,7 +39,7 @@ def load_env(arena_env_args: list[str], job_name: str):
     if hasattr(env_cfg, "recorders") and env_cfg.recorders is not None:
         env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
 
-    env = arena_builder.make_registered(env_cfg)
+    env = arena_builder.make_registered(env_cfg, render_mode=render_mode)
     # Don't reset here - rollout_policy() will reset the env. Every reset triggers a new episode, initializing recorder & creating a new hdf5 entry.
     return env
 
@@ -113,12 +114,16 @@ def main():
 
         job_manager.print_jobs_info()
 
+        if args_cli.video:
+            os.makedirs(args_cli.video_dir, exist_ok=True)
+            print(f"[INFO] Video recording enabled. Videos will be saved to: {args_cli.video_dir}")
+
         for job in job_manager:
             if job is not None:
                 env = None
                 try:
-                    # Modules reloading first, otherwise 2 instances of same class are created (e.g. Enum)
-                    env = load_env(job.arena_env_args, job.name)
+                    render_mode = "rgb_array" if args_cli.video else None
+                    env = load_env(job.arena_env_args, job.name, render_mode=render_mode)
 
                     policy = get_policy_from_job(job)
 
@@ -129,6 +134,21 @@ def main():
                             job.num_steps = policy.length()
                         else:
                             job.num_steps = args_cli.num_steps
+
+                    if args_cli.video:
+                        if job.num_steps is not None:
+                            video_length = job.num_steps
+                        else:
+                            video_length = job.num_episodes * env.unwrapped.max_episode_length
+                        video_kwargs = {
+                            "video_folder": os.path.join(args_cli.video_dir, job.name),
+                            "step_trigger": lambda step: step == 0,
+                            "video_length": video_length,
+                            "disable_logger": True,
+                        }
+                        print(f"[INFO] Recording video for job '{job.name}' -> {video_kwargs['video_folder']}")
+                        env = RecordVideo(env, **video_kwargs)
+
                     metrics = rollout_policy(env, policy, num_steps=job.num_steps, num_episodes=job.num_episodes)
 
                     job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)
@@ -138,10 +158,11 @@ def main():
                         metrics_logger.append_job_metrics(job.name, metrics)
 
                 except Exception as e:
-                    # continue with the next job even if one fails
                     job_manager.complete_job(job, metrics={}, status=Status.FAILED)
                     print(f"Job {job.name} failed with error: {e}")
                     print(f"Traceback: {traceback.format_exc()}")
+                    if not args_cli.continue_on_error:
+                        raise
 
                 finally:
                     # Only stop env if it was successfully created

isaaclab_arena/evaluation/eval_runner_cli.py

@@ -14,3 +14,16 @@ def add_eval_runner_arguments(parser: argparse.ArgumentParser) -> None:
         default="isaaclab_arena_environments/eval_jobs_configs/zero_action_jobs_config.json",
         help="Path to the eval jobs config file.",
     )
+    parser.add_argument(
+        "--continue_on_error",
+        action="store_true",
+        default=False,
+        help="Continue evaluation with remaining jobs when a job fails instead of stopping immediately.",
+    )
+    parser.add_argument("--video", action="store_true", default=False, help="Record videos for each eval job.")
+    parser.add_argument(
+        "--video_dir",
+        type=str,
+        default="/eval/videos",
+        help="Root directory for recorded videos. Each job gets a subdirectory.",
+    )

isaaclab_arena/evaluation/policy_runner.py

@@ -63,11 +63,13 @@ def rollout_policy(
     assert num_steps is not None or num_episodes is not None, "Either num_steps or num_episodes must be provided"
     assert num_steps is None or num_episodes is None, "Only one of num_steps or num_episodes must be provided"
 
+    pbar = None
     try:
         obs, _ = env.reset()
         policy.reset()
         # set task description (could be None) from the task being evaluated
-        policy.set_task_description(env.cfg.isaaclab_arena_env.task.get_task_description())
+        # Use unwrapped to reach the base env through any gym wrappers (e.g. RecordVideo)
+        policy.set_task_description(env.unwrapped.cfg.isaaclab_arena_env.task.get_task_description())
 
         # Setup progress bar based on num_steps or num_episodes
         if num_steps is not None:
@@ -107,16 +109,17 @@ def rollout_policy(
         pbar.close()
 
     except Exception as e:
-        pbar.close()
+        if pbar is not None:
+            pbar.close()
         raise RuntimeError(f"Error rolling out policy: {e}")
 
     else:
         # only compute metrics if env has metrics registered
-        if hasattr(env.cfg, "metrics"):
+        if hasattr(env.unwrapped.cfg, "metrics"):
             # NOTE(xinjieyao, 2025-10-07): lazy import to prevent app stalling caused by omni.kit
             from isaaclab_arena.metrics.metrics import compute_metrics
 
-            metrics = compute_metrics(env)
+            metrics = compute_metrics(env.unwrapped)
             return metrics
         return None
 

isaaclab_arena/policy/pi0_droid_client.py

@@ -0,0 +1,138 @@
+# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pi0 client policy for DROID environments.
+
+
+# TODO(cvolk)
+/isaac-sim/python.sh -m isaaclab_arena.evaluation.policy_runner   --policy_type isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient   --enable_cameras   --num_episodes 10   --remote_host localhost   --remote_port 8000   droid_pick_and_place_srl   --embodiment droid_rel_joint_pos
+
+
+Thin client that connects directly to an openpi WebSocket server.
+No custom server wrapper needed — uses the upstream pi0 server as-is.
+
+pi0_fast_droid outputs joint *velocities* (7) + gripper position (1),
+clipped to [-1, 1]. Use with ``droid_rel_joint_pos`` embodiment which
+applies the first 7 dims as deltas from the current joint position.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gymnasium as gym
+import numpy as np
+import torch
+from dataclasses import dataclass
+from gymnasium.spaces.dict import Dict as GymSpacesDict
+from typing import Any
+
+from openpi_client import image_tools, websocket_client_policy
+
+from isaaclab_arena.policy.policy_base import PolicyBase
+
+
+@dataclass
+class Pi0DroidClientConfig:
+    remote_host: str = "localhost"
+    remote_port: int = 8000
+    open_loop_horizon: int = 8
+    image_size: int = 224
+    external_camera_key: str = "external_camera_rgb"
+    wrist_camera_key: str = "wrist_camera_rgb"
+
+
+class Pi0DroidClient(PolicyBase):
+    """Client-side policy that talks directly to an openpi WebSocket server.
+
+    Designed for DROID embodiment with joint-position action space.
+    Handles action chunking, image resizing, and gripper binarization locally.
+    """
+
+    config_class = Pi0DroidClientConfig
+
+    def __init__(self, config: Pi0DroidClientConfig) -> None:
+        super().__init__(config=config)
+        self.cfg = config
+
+        print(f"[Pi0DroidClient] Connecting to openpi server at {config.remote_host}:{config.remote_port}...")
+        self.client = websocket_client_policy.WebsocketClientPolicy(config.remote_host, config.remote_port)
+        print("[Pi0DroidClient] Server ready.")
+
+        self._chunk: np.ndarray | None = None
+        self._chunk_idx: int = 0
+        self.task_description: str | None = None
+
+    def get_action(self, env: gym.Env, observation: GymSpacesDict) -> torch.Tensor:
+        if self._chunk is None or self._chunk_idx >= self.cfg.open_loop_horizon:
+            self._chunk_idx = 0
+            self._chunk = self._request_chunk(observation)
+
+        action_np = self._chunk[self._chunk_idx].copy()
+        self._chunk_idx += 1
+
+        # pi0 outputs joint velocities clipped to [-1, 1] for the arm,
+        # and a gripper position for the last dim.
+        action_np[:7] = np.clip(action_np[:7], -1.0, 1.0)
+        action_np[-1] = 1.0 if action_np[-1] > 0.5 else 0.0
+        return torch.tensor(action_np, dtype=torch.float32, device="cuda").unsqueeze(0)
+
+    def reset(self, env_ids: torch.Tensor | None = None) -> None:
+        self._chunk = None
+        self._chunk_idx = 0
+
+    def set_task_description(self, task_description: str | None) -> str:
+        self.task_description = task_description
+        return self.task_description or ""
+
+    def _request_chunk(self, observation: dict[str, Any]) -> np.ndarray:
+        obs = self._extract_observation(observation)
+        sz = self.cfg.image_size
+
+        request = {
+            "observation/exterior_image_1_left": image_tools.resize_with_pad(obs["external_image"], sz, sz),
+            "observation/wrist_image_left": image_tools.resize_with_pad(obs["wrist_image"], sz, sz),
+            "observation/joint_position": obs["joint_position"],
+            "observation/gripper_position": obs["gripper_position"],
+            "prompt": self.task_description or "",
+        }
+
+        response = self.client.infer(request)
+        return response["actions"]
+
+    def _extract_observation(self, obs_dict: dict[str, Any], env_id: int = 0) -> dict[str, np.ndarray]:
+        ext_image = obs_dict["camera_obs"][self.cfg.external_camera_key][env_id].detach().cpu().numpy()
+        wrist_image = obs_dict["camera_obs"][self.cfg.wrist_camera_key][env_id].detach().cpu().numpy()
+        joint_pos = obs_dict["policy"]["joint_pos"][env_id].detach().cpu().numpy()
+        gripper_pos = obs_dict["policy"]["gripper_pos"][env_id].detach().cpu().numpy()
+
+        return {
+            "external_image": ext_image,
+            "wrist_image": wrist_image,
+            "joint_position": joint_pos,
+            "gripper_position": gripper_pos,
+        }
+
+    @staticmethod
+    def add_args_to_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        group = parser.add_argument_group("Pi0 DROID Client")
+        group.add_argument("--remote_host", type=str, default="localhost")
+        group.add_argument("--remote_port", type=int, default=8000)
+        group.add_argument("--open_loop_horizon", type=int, default=8)
+        group.add_argument("--image_size", type=int, default=224)
+        group.add_argument("--external_camera_key", type=str, default="external_camera_rgb")
+        group.add_argument("--wrist_camera_key", type=str, default="wrist_camera_rgb")
+        return parser
+
+    @staticmethod
+    def from_args(args: argparse.Namespace) -> Pi0DroidClient:
+        config = Pi0DroidClientConfig(
+            remote_host=args.remote_host,
+            remote_port=args.remote_port,
+            open_loop_horizon=args.open_loop_horizon,
+            image_size=args.image_size,
+            external_camera_key=args.external_camera_key,
+            wrist_camera_key=args.wrist_camera_key,
+        )
+        return Pi0DroidClient(config)

isaaclab_arena_environments/droid_pick_and_place_srl_environment.py

@@ -25,7 +25,7 @@ def get_env(self, args_cli: argparse.Namespace):  # -> IsaacLabArenaEnvironment:
         from isaaclab_arena.assets.object_base import ObjectType
         from isaaclab_arena.assets.object_reference import ObjectReference
         from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment
-        from isaaclab_arena.relations.relations import IsAnchor, On
+        from isaaclab_arena.relations.relations import AtPosition, IsAnchor, On, RandomAroundSolution
         from isaaclab_arena.scene.scene import Scene
         from isaaclab_arena.tasks.pick_and_place_task import PickAndPlaceTask
 
@@ -43,6 +43,8 @@ def get_env(self, args_cli: argparse.Namespace):  # -> IsaacLabArenaEnvironment:
         pick_up_object.add_relation(On(table_reference))
         destination_location = self.asset_registry.get_asset_by_name(args_cli.destination_location)()
         destination_location.add_relation(On(table_reference))
+        destination_location.add_relation(AtPosition(x=0.5, y=0.0))
+        destination_location.add_relation(RandomAroundSolution(x_half_m=0.2, y_half_m=0.25))
 
         light = self.asset_registry.get_asset_by_name("light")(
             spawner_cfg=sim_utils.DomeLightCfg(intensity=500.0),