diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index 772923ae..a662d644 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/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 diff --git a/isaaclab_arena/evaluation/eval_runner.py b/isaaclab_arena/evaluation/eval_runner.py index d3f3c294..09a9f8f0 100644 --- a/isaaclab_arena/evaluation/eval_runner.py +++ b/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 diff --git a/isaaclab_arena/evaluation/eval_runner_cli.py b/isaaclab_arena/evaluation/eval_runner_cli.py index 74307da0..909e31a4 100644 --- a/isaaclab_arena/evaluation/eval_runner_cli.py +++ b/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.", + ) diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py index dec7e2ea..6d7f1de7 100644 --- a/isaaclab_arena/evaluation/policy_runner.py +++ b/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 diff --git a/isaaclab_arena/policy/pi0_droid_client.py b/isaaclab_arena/policy/pi0_droid_client.py new file mode 100644 index 00000000..52a58cfc --- /dev/null +++ b/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) diff --git a/isaaclab_arena_environments/droid_pick_and_place_srl_environment.py b/isaaclab_arena_environments/droid_pick_and_place_srl_environment.py index 1980c0fb..9b0fed0e 100644 --- a/isaaclab_arena_environments/droid_pick_and_place_srl_environment.py +++ b/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), diff --git a/isaaclab_arena_environments/eval_jobs_configs/droid_pnp_srl_openpi_jobs_config.json b/isaaclab_arena_environments/eval_jobs_configs/droid_pnp_srl_openpi_jobs_config.json new file mode 100644 index 00000000..127b8024 --- /dev/null +++ b/isaaclab_arena_environments/eval_jobs_configs/droid_pnp_srl_openpi_jobs_config.json @@ -0,0 +1,115 @@ +{ + "jobs": [ + { + "name": "droid_pnp_srl_openpi_billiard_hall", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "hdr": "billiard_hall_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + }, + { + "name": "droid_pnp_srl_openpi_billiard_hall_wooden_bowl", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "hdr": "billiard_hall_robolab", + "destination_location": "bowl_ycb_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + }, + { + "name": "droid_pnp_srl_openpi_blue_block", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "pick_up_object": "blue_block_basic_robolab", + "destination_location": "bowl_ycb_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + }, + { + "name": "droid_pnp_srl_openpi_alphabet_soup_can", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "pick_up_object": "alphabet_soup_can_hope_robolab", + "destination_location": "bowl_ycb_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + }, + { + "name": "droid_pnp_srl_openpi_orange", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "pick_up_object": "orange_01_fruits_veggies_robolab", + "destination_location": "bowl_ycb_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + }, + { + "name": "droid_pnp_srl_openpi_lemon", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "pick_up_object": "lemon_01_fruits_veggies_robolab", + "destination_location": "bowl_ycb_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + }, + { + "name": "droid_pnp_srl_openpi_tomato_sauce_can", + "arena_env_args": { + "enable_cameras": true, + "environment": "droid_pick_and_place_srl", + "embodiment": "droid_rel_joint_pos", + "pick_up_object": "tomato_sauce_can_hot3d_robolab", + "destination_location": "bowl_ycb_robolab" + }, + "num_episodes": 3, + "policy_type": "isaaclab_arena.policy.pi0_droid_client.Pi0DroidClient", + "policy_config_dict": { + "remote_host": "localhost", + "remote_port": 8000 + } + } + ] +}