diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7a023bb2b4..e588fd9fc2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -308,16 +308,13 @@ jobs: - name: Run in-process PhysX tests without cameras run: | isaaclab.sh -p -m pytest -sv --durations=0 -m "not with_cameras and not with_subprocess and not with_newton" \ - isaaclab_arena/tests/ + isaaclab_arena/tests/ \ + isaaclab_arena_vlm_agent/tests/ - name: Run in-process PhysX tests with cameras run: | isaaclab.sh -p -m pytest -sv --durations=0 -m "with_cameras and not with_subprocess and not with_newton" \ - isaaclab_arena/tests/ - - - name: Run VLM agent policy tests - run: | - isaaclab.sh -p -m pytest -sv --durations=0 \ + isaaclab_arena/tests/ \ isaaclab_arena_vlm_agent/tests/ - name: Run GR00T policy client & data tests (lightweight gr00t deps only) @@ -364,7 +361,8 @@ jobs: run: | ISAACLAB_ARENA_SUBPROCESS_TIMEOUT=900 \ isaaclab.sh -p -m pytest -sv --durations=0 -m "with_subprocess and not with_newton" \ - isaaclab_arena/tests/ + isaaclab_arena/tests/ \ + isaaclab_arena_vlm_agent/tests/ test_gr00t_closedloop_e2e: diff --git a/.github/workflows/uv-ci.yml b/.github/workflows/uv-ci.yml index 7b3ecd11c0..433f8fa7c3 100644 --- a/.github/workflows/uv-ci.yml +++ b/.github/workflows/uv-ci.yml @@ -170,14 +170,16 @@ jobs: run: | source .venv/bin/activate pytest -sv --durations=0 -m "not with_cameras and not with_subprocess and not with_newton" \ - isaaclab_arena/tests/ + isaaclab_arena/tests/ \ + isaaclab_arena_vlm_agent/tests/ - name: Run in-process PhysX tests with cameras shell: bash run: | source .venv/bin/activate pytest -sv --durations=0 -m "with_cameras and not with_subprocess and not with_newton" \ - isaaclab_arena/tests/ + isaaclab_arena/tests/ \ + isaaclab_arena_vlm_agent/tests/ - name: Run subprocess-spawning Newton tests shell: bash @@ -193,13 +195,7 @@ jobs: source .venv/bin/activate ISAACLAB_ARENA_SUBPROCESS_TIMEOUT=900 \ pytest -sv --durations=0 -m "with_subprocess and not with_newton" \ - isaaclab_arena/tests/ - - - name: Run VLM agent policy tests - shell: bash - run: | - source .venv/bin/activate - pytest -sv --durations=0 \ + isaaclab_arena/tests/ \ isaaclab_arena_vlm_agent/tests/ - name: Run GR00T remote policy client tests diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index cb8637b033..ce3abe23df 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -291,6 +291,27 @@ def _configure_newton_diff_ik(self) -> None: ) +@register_asset +class DroidAbsoluteDifferentialIKEmbodiment(DroidDifferentialIKEmbodiment): + """DROID differential IK accepting unscaled root-relative XYZ/XYZW pose targets.""" + + registry_name = "droid_abs_ik" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._set_absolute_pose_actions() + + def _configure_physics_backend(self, backend: PhysicsBackend) -> None: + """Apply backend IK settings, then accept unscaled absolute pose targets.""" + super()._configure_physics_backend(backend) + self._set_absolute_pose_actions() + + def _set_absolute_pose_actions(self) -> None: + """Keep the stored action config in absolute pose mode.""" + self.action_config.arm_action.controller.use_relative_mode = False + self.action_config.arm_action.scale = 1.0 + + @register_asset class DroidRelativeJointPositionEmbodiment(DroidEmbodimentBase): """Embodiment for the DROID setup with relative joint position action controller.""" @@ -569,6 +590,7 @@ class DroidCameraCfg(ArenaCameraCfg): """Configuration for cameras. DROID cameras are mounted with pre-set poses.""" external_camera: CameraCfg = CameraCfg( + update_latest_camera_pose=True, prim_path="{ENV_REGEX_NS}/Robot/panda_link0/external_camera", height=720, width=1280, @@ -582,6 +604,7 @@ class DroidCameraCfg(ArenaCameraCfg): offset=CameraCfg.OffsetCfg(pos=(0.05, 0.57, 0.66), rot=(-0.195, 0.399, 0.805, -0.393), convention="opengl"), ) external_camera_2: CameraCfg = CameraCfg( + update_latest_camera_pose=True, prim_path="{ENV_REGEX_NS}/Robot/panda_link0/external_camera_2", height=720, width=1280, @@ -595,6 +618,7 @@ class DroidCameraCfg(ArenaCameraCfg): offset=CameraCfg.OffsetCfg(pos=(0.05, -0.57, 0.66), rot=(0.399, -0.195, -0.393, 0.805), convention="opengl"), ) wrist_camera: CameraCfg = CameraCfg( + update_latest_camera_pose=True, prim_path="{ENV_REGEX_NS}/Robot/Gripper/Robotiq_2F_85/base_link/wrist_camera", height=720, width=1280, diff --git a/isaaclab_arena/embodiments/droid/observations.py b/isaaclab_arena/embodiments/droid/observations.py index 480eb4608e..1fc7c88df3 100644 --- a/isaaclab_arena/embodiments/droid/observations.py +++ b/isaaclab_arena/embodiments/droid/observations.py @@ -54,7 +54,7 @@ def ee_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("r def ee_quat(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: - """Returns the end effector orientation as quaternion (w, x, y, z) in the world frame.""" + """Returns the end effector orientation as quaternion (x, y, z, w) in the world frame.""" robot = env.scene[asset_cfg.name] body_idx = robot.data.body_names.index("base_link") # Robotiq gripper base link return wp.to_torch(robot.data.body_quat_w)[:, body_idx, :] diff --git a/isaaclab_arena/inference/backend.py b/isaaclab_arena/inference/backend.py index 767a409834..31c631f0cb 100644 --- a/isaaclab_arena/inference/backend.py +++ b/isaaclab_arena/inference/backend.py @@ -188,6 +188,10 @@ def client(self) -> OpenAI: """OpenAI-compatible client used for completion requests.""" return self._client + def close(self) -> None: + """Close the owned inference client.""" + self._client.close() + def run_json(self, request: StructuredOutputRequest) -> dict[str, Any]: """Call a JSON-schema structured-output endpoint and parse the response as JSON. diff --git a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_chunk_experiment.yaml b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_chunk_experiment.yaml new file mode 100644 index 0000000000..4c4dd234da --- /dev/null +++ b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_chunk_experiment.yaml @@ -0,0 +1,35 @@ +# Copyright (c) 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 + +# Run inside the Arena runtime with NV_API_KEY set: +# /isaac-sim/python.sh isaaclab_arena/evaluation/experiment_runner.py --viz none \ +# --experiment_config isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_chunk_experiment.yaml +# Override endpoint/model/base_url as needed. Credentials stay in the environment. +# The droid preset supplies absolute-pose adapters, Robotiq measurements, and strategy-specific prompts. +# Set policy.system_prompt to override its default prompt. +runs: + droid_pnp_vlm_chunk: + environment: + type: pick_and_place_maple_table + enable_cameras: true + embodiment: droid_abs_ik + pick_up_object: rubiks_cube_hot3d_robolab + destination_location: bowl_ycb_robolab + hdr: home_office_robolab + environment_builder: + num_envs: 1 + language_instruction: >- + Pick up the Rubik's cube and place it in the bowl. + Approach with the gripper open, close around the object, lift it, + move above the destination, and open to release. + policy: + type: isaaclab_arena_vlm_agent.vlm_agent_policy.VLMChunkPolicy + endpoint: internal + model: openai/openai/gpt-6-astra + vlm_embodiment_adapter: droid + observation_modality: image_state + chunk_size: 3 + rollout_limit: + num_episodes: 1 diff --git a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_goal_experiment.yaml b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_goal_experiment.yaml new file mode 100644 index 0000000000..4460c03d3d --- /dev/null +++ b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_goal_experiment.yaml @@ -0,0 +1,35 @@ +# Copyright (c) 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 + +# Run inside the Arena runtime with NV_API_KEY set: +# /isaac-sim/python.sh isaaclab_arena/evaluation/experiment_runner.py --viz none \ +# --experiment_config isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_goal_experiment.yaml +# Override endpoint/model/base_url as needed. Credentials stay in the environment. +# The droid preset supplies absolute-pose adapters, Robotiq measurements, and strategy-specific prompts. +# Set policy.system_prompt to override its default prompt. +runs: + droid_pnp_vlm_goal: + environment: + type: pick_and_place_maple_table + enable_cameras: true + embodiment: droid_abs_ik + pick_up_object: rubiks_cube_hot3d_robolab + destination_location: bowl_ycb_robolab + hdr: home_office_robolab + environment_builder: + num_envs: 1 + language_instruction: >- + Pick up the Rubik's cube and place it in the bowl. + Approach with the gripper open, close around the object, lift it, + move above the destination, and open to release. + policy: + type: isaaclab_arena_vlm_agent.vlm_agent_policy.VLMGoalPolicy + endpoint: internal + model: openai/openai/gpt-6-astra + vlm_embodiment_adapter: droid + observation_modality: image_state + max_steps_per_goal: 30 + rollout_limit: + num_episodes: 1 diff --git a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_step_experiment.yaml b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_step_experiment.yaml new file mode 100644 index 0000000000..05b6f8d75d --- /dev/null +++ b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_step_experiment.yaml @@ -0,0 +1,34 @@ +# Copyright (c) 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 + +# Run inside the Arena runtime with NV_API_KEY set: +# /isaac-sim/python.sh isaaclab_arena/evaluation/experiment_runner.py --viz none \ +# --experiment_config isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_step_experiment.yaml +# Override endpoint/model/base_url as needed. Credentials stay in the environment. +# The droid preset supplies absolute-pose adapters, Robotiq measurements, and strategy-specific prompts. +# Set policy.system_prompt to override its default prompt. +runs: + droid_pnp_vlm_step: + environment: + type: pick_and_place_maple_table + enable_cameras: true + embodiment: droid_abs_ik + pick_up_object: rubiks_cube_hot3d_robolab + destination_location: bowl_ycb_robolab + hdr: home_office_robolab + environment_builder: + num_envs: 1 + language_instruction: >- + Pick up the Rubik's cube and place it in the bowl. + Approach with the gripper open, close around the object, lift it, + move above the destination, and open to release. + policy: + type: isaaclab_arena_vlm_agent.vlm_agent_policy.VLMStepPolicy + endpoint: internal + model: openai/openai/gpt-6-astra + vlm_embodiment_adapter: droid + observation_modality: image_state + rollout_limit: + num_episodes: 1 diff --git a/isaaclab_arena_vlm_agent/absolute_action_adapter.py b/isaaclab_arena_vlm_agent/absolute_action_adapter.py index 14301b4112..8399addaae 100644 --- a/isaaclab_arena_vlm_agent/absolute_action_adapter.py +++ b/isaaclab_arena_vlm_agent/absolute_action_adapter.py @@ -7,54 +7,62 @@ import numpy as np -from isaaclab_arena_vlm_agent.action_adapter import ChunkActionAdapter, GoalActionAdapter +from isaaclab_arena_vlm_agent.action_adapter import AgentActionAdapter, ChunkActionAdapter, GoalActionAdapter from isaaclab_arena_vlm_agent.command_indices import AbsoluteCommandIndices, read_absolute_command_indices from isaaclab_arena_vlm_agent.commands import Command, CommandChunk from isaaclab_arena_vlm_agent.motion_generator import MotionGenerator from isaaclab_arena_vlm_agent.response_errors import InvalidAgentResponseError -class AbsoluteChunkActionAdapter(ChunkActionAdapter): - """Decode absolute command chunks; use chunk_size=1 for one prediction per step.""" +class _AbsoluteCommandAdapter(AgentActionAdapter): + """Decode a single absolute command in native controller order.""" - response_model = CommandChunk + response_model = Command - def __init__(self, chunk_size: int): - super().__init__(chunk_size) + def __init__(self): self._command_indices = None def is_robot_action_space_supported(self, env) -> bool: - """Resolve command indices from the active controller before policy inference.""" + """Read absolute command indices before policy inference.""" self._command_indices = read_absolute_command_indices(env) return self._command_indices is not None + def decode_command(self, response: Command, arena_world, env_id: int) -> np.ndarray: + """Decode and validate one absolute target.""" + return _decode_absolute_command(response, self._command_indices) + + +class AbsoluteStepActionAdapter(_AbsoluteCommandAdapter): + """Decode a single absolute command in native controller order.""" + + +class AbsoluteChunkActionAdapter(_AbsoluteCommandAdapter, ChunkActionAdapter): + """Decode absolute command chunks in native controller order.""" + + response_model = CommandChunk + + def __init__(self, chunk_size: int): + _AbsoluteCommandAdapter.__init__(self) + ChunkActionAdapter.__init__(self, chunk_size) + def decode_command(self, response: CommandChunk, arena_world, env_id: int) -> np.ndarray: """Validate and stack absolute commands in execution order.""" if len(response.commands) != self.chunk_size: raise InvalidAgentResponseError(f"Expected {self.chunk_size} commands, got {len(response.commands)}") commands = [] for command in response.commands: - commands.append(_decode_absolute_command(command, self._command_indices)) + commands.append(super().decode_command(command, arena_world, env_id)) return np.stack(commands) -class AbsoluteGoalActionAdapter(GoalActionAdapter): +class AbsoluteGoalActionAdapter(_AbsoluteCommandAdapter, GoalActionAdapter): """Decode absolute targets for an explicitly supplied motion generator.""" response_model = Command def __init__(self, motion_generator: MotionGenerator): - super().__init__(motion_generator) - self._command_indices = None - - def is_robot_action_space_supported(self, env) -> bool: - """Resolve command indices from the active controller before policy inference.""" - self._command_indices = read_absolute_command_indices(env) - return self._command_indices is not None - - def decode_command(self, response: Command, arena_world, env_id: int) -> np.ndarray: - """Decode a validated response into an absolute target for the motion generator.""" - return _decode_absolute_command(response, self._command_indices) + _AbsoluteCommandAdapter.__init__(self) + GoalActionAdapter.__init__(self, motion_generator) def _decode_absolute_command(response: Command, indices: AbsoluteCommandIndices | None) -> np.ndarray: diff --git a/isaaclab_arena_vlm_agent/droid_adapter.py b/isaaclab_arena_vlm_agent/droid_adapter.py new file mode 100644 index 0000000000..0892b09365 --- /dev/null +++ b/isaaclab_arena_vlm_agent/droid_adapter.py @@ -0,0 +1,97 @@ +# Copyright (c) 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 + +"""DROID absolute-pose action, motion-generation, and prompt defaults.""" + +from isaaclab_arena_vlm_agent.absolute_action_adapter import ( + AbsoluteChunkActionAdapter, + AbsoluteGoalActionAdapter, + AbsoluteStepActionAdapter, +) +from isaaclab_arena_vlm_agent.action_adapter import AgentActionAdapter +from isaaclab_arena_vlm_agent.command_indices import AbsoluteCommandIndices, GripperCommandIndex, PoseCommandIndices +from isaaclab_arena_vlm_agent.pose_target_motion_generator import PoseTargetMotionGenerator + + +class _DroidPoseActionAdapter: + """Require the native pose and gripper layout described by the DROID prompt.""" + + def is_robot_action_space_supported(self, env) -> bool: + return super().is_robot_action_space_supported(env) and self._command_indices == AbsoluteCommandIndices( + action_dim=8, + poses=(PoseCommandIndices(slice(0, 7), "robot", "base_link"),), + joints=(), + grippers=(GripperCommandIndex(7, "gripper_action", 0.0, 1.0),), + ) + + +class _DroidStepActionAdapter(_DroidPoseActionAdapter, AbsoluteStepActionAdapter): + """Decode a DROID pose command.""" + + +class _DroidChunkActionAdapter(_DroidPoseActionAdapter, AbsoluteChunkActionAdapter): + """Decode a chunk of DROID pose commands.""" + + +class _DroidGoalActionAdapter(_DroidPoseActionAdapter, AbsoluteGoalActionAdapter): + """Decode a DROID pose goal.""" + + +class DroidVLMAdapter: + """Provide VLM defaults for the droid_abs_ik embodiment.""" + + def make_step_action_adapter(self) -> AbsoluteStepActionAdapter: + """Create a decoder for one absolute pose command.""" + return _DroidStepActionAdapter() + + def make_chunk_action_adapter(self, chunk_size: int) -> AbsoluteChunkActionAdapter: + """Create a decoder for the requested number of absolute pose commands.""" + return _DroidChunkActionAdapter(chunk_size) + + def make_goal_action_adapter( + self, position_tolerance_m: float, rotation_tolerance_rad: float + ) -> AbsoluteGoalActionAdapter: + """Create pose-goal execution with DROID's Robotiq gripper measurements.""" + generator = PoseTargetMotionGenerator( + grippers={"gripper_action": _make_gripper()}, + position_tolerance_m=position_tolerance_m, + rotation_tolerance_rad=rotation_tolerance_rad, + ) + return _DroidGoalActionAdapter(generator) + + def get_system_prompt(self, action_adapter: AgentActionAdapter) -> str: + """Describe the DROID pose contract and the selected response schema.""" + if isinstance(action_adapter, AbsoluteChunkActionAdapter): + response = ( + f"Return a JSON object with a commands list of exactly {action_adapter.chunk_size} commands, " + "executed one per simulation step. Each command has a values list of eight numbers: " + ) + else: + assert isinstance( + action_adapter, (AbsoluteStepActionAdapter, AbsoluteGoalActionAdapter) + ), "Provide system_prompt when injecting a custom action adapter" + response = "Return one command as a JSON object with a values list of eight numbers: " + return ( + "Control DROID to complete the supplied task using the current camera images and robot state. " + + response + + "absolute XYZ in meters and a unit XYZW quaternion for the gripper base_link body, expressed in " + "the robot root frame, followed by the gripper command (0 opens, 1 closes). Read the current " + "camera_geometry for each named image: image_size_px is [width, height], intrinsic_matrix is K " + "in the transmitted image's pixels, and camera poses map optical coordinates (+X right, +Y down, " + "+Z forward) into the named world or robot root frame. For a camera point p_C, " + "p_R = rotate(quaternion_root_xyzw, p_C) + position_root_m; pixels project as K p_C / p_C.z. " + "All quaternions in input and output use XYZW. policy.eef_pos is in world meters and " + "policy.eef_quat is the gripper's world orientation in XYZW. Read the current " + "commanded-body pose from command_body_poses.arm_action when state is provided. This body is not " + "the finger midpoint. Predict nearby waypoints using small pose changes. " + "Follow the response schema exactly, with no explanatory text." + ) + + +def _make_gripper(): + """Create the measurements used by the DROID embodiment after simulation startup.""" + from isaaclab_arena.embodiments.gripper import RobotiqGripper + + return RobotiqGripper() diff --git a/isaaclab_arena_vlm_agent/image_utils.py b/isaaclab_arena_vlm_agent/image_utils.py index 1cf7eb6921..2ca5b6eb0c 100644 --- a/isaaclab_arena_vlm_agent/image_utils.py +++ b/isaaclab_arena_vlm_agent/image_utils.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Image encoding utilities.""" +"""Image encoding and camera calibration utilities.""" import base64 import io @@ -34,7 +34,7 @@ def extract_rgb_images_from_camera_observations( return images -def image_to_jpeg_data_url(image: np.ndarray, max_width_or_height: int) -> str: +def image_to_jpeg_data_url(image: np.ndarray, max_width_or_height: int) -> tuple[str, tuple[int, int]]: """Encode an RGB or RGBA image as JPEG, preserving aspect ratio without upscaling. Args: @@ -42,7 +42,7 @@ def image_to_jpeg_data_url(image: np.ndarray, max_width_or_height: int) -> str: max_width_or_height: Maximum encoded image width or height, in pixels. Returns: - A JPEG data URL for an image-capable chat-completion endpoint. + A JPEG data URL and its actual (width, height) in pixels. """ assert max_width_or_height > 0, "max_width_or_height must be positive" assert image.ndim == 3 and image.shape[-1] in (3, 4), f"Unexpected image shape {image.shape}" @@ -51,4 +51,19 @@ def image_to_jpeg_data_url(image: np.ndarray, max_width_or_height: int) -> str: frame.thumbnail((max_width_or_height, max_width_or_height)) buffer = io.BytesIO() frame.save(buffer, format="JPEG", quality=85) - return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode("ascii") + return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode("ascii"), frame.size + + +def resize_camera_geometry(geometry: dict, original_size: tuple[int, int], resized_size: tuple[int, int]) -> None: + """Update camera intrinsics and image dimensions in place after resizing. + + Args: + geometry: Camera calibration containing an intrinsic_matrix and image_size_px. + original_size: Original image (width, height) in pixels. + resized_size: Actual resized image (width, height) in pixels. + """ + intrinsics = np.array(geometry["intrinsic_matrix"], dtype=float) + intrinsics[0] *= resized_size[0] / original_size[0] + intrinsics[1] *= resized_size[1] / original_size[1] + geometry["intrinsic_matrix"] = intrinsics.tolist() + geometry["image_size_px"] = list(resized_size) diff --git a/isaaclab_arena_vlm_agent/observation_adapter.py b/isaaclab_arena_vlm_agent/observation_adapter.py index 1ba313caa6..45825efc93 100644 --- a/isaaclab_arena_vlm_agent/observation_adapter.py +++ b/isaaclab_arena_vlm_agent/observation_adapter.py @@ -9,13 +9,11 @@ import numpy as np from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from scipy.spatial.transform import Rotation +from typing import Any from isaaclab_arena_vlm_agent.image_utils import extract_rgb_images_from_camera_observations -if TYPE_CHECKING: - from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase - POLICY_OBSERVATION_GROUP = "policy" """Observation group containing the embodiment's named state terms.""" @@ -92,36 +90,93 @@ def extract_vlm_input_from_observation(self, env, observation, env_id: int): class ImageObservationAdapter(VLMObservationAdapter): - """Include the configured current RGB camera views for one environment.""" + """Include RGB views and calibration from the live environment.""" - def __init__(self, embodiment: EmbodimentBase): - observation_cfg = embodiment.get_observation_cfg() - camera_cfg = getattr(observation_cfg, CAMERA_OBSERVATION_GROUP, None) + def extract_vlm_input_from_observation(self, env, observation, env_id: int): + """Return RGB views and optical-frame camera calibration for one environment.""" + base = env.unwrapped + camera_cfg = getattr(base.cfg.observations, CAMERA_OBSERVATION_GROUP, None) assert camera_cfg is not None, "RGB modality requires camera observations; enable the embodiment's cameras" - self.camera_keys = tuple( + camera_keys = tuple( name for name, term in vars(camera_cfg).items() if hasattr(term, "func") and term.params.get("data_type") == "rgb" ) - assert self.camera_keys, "RGB modality requires at least one configured RGB observation" - - def extract_vlm_input_from_observation(self, env, observation, env_id: int): - """Return configured RGB views for the selected environment.""" - images = extract_rgb_images_from_camera_observations( - observation[CAMERA_OBSERVATION_GROUP], self.camera_keys, env_id - ) - return images, {} + assert camera_keys, "RGB modality requires at least one configured RGB observation" + images = extract_rgb_images_from_camera_observations(observation[CAMERA_OBSERVATION_GROUP], camera_keys, env_id) + return images, {"camera_geometry": _extract_camera_geometry(env, images, env_id)} class ImageStateObservationAdapter(VLMObservationAdapter): """Include robot state followed by current RGB camera views.""" - def __init__(self, embodiment: EmbodimentBase): + def __init__(self): self.state_adapter = StateObservationAdapter() - self.image_adapter = ImageObservationAdapter(embodiment) + self.image_adapter = ImageObservationAdapter() def extract_vlm_input_from_observation(self, env, observation, env_id: int): """Return state and RGB views for the selected environment.""" _, state = self.state_adapter.extract_vlm_input_from_observation(env, observation, env_id) - images, _ = self.image_adapter.extract_vlm_input_from_observation(env, observation, env_id) - return images, state + images, geometry = self.image_adapter.extract_vlm_input_from_observation(env, observation, env_id) + return images, {**state, **geometry} + + +def _extract_camera_geometry(env, images: dict[str, np.ndarray], env_id: int) -> dict[str, dict[str, Any]]: + """Read camera intrinsics and optical-frame poses for the selected RGB views. + + Args: + env: Arena environment, optionally gym-wrapped, with live camera sensors. + images: HWC RGB images indexed by their observation term names. + env_id: Environment index whose camera calibration is requested. + + Returns: + Calibration by observation term, with world and robot-root poses using XYZW quaternions. + """ + base = env.unwrapped + camera_cfg = getattr(base.cfg.observations, CAMERA_OBSERVATION_GROUP) + roots = {} + for name in base.action_manager.active_terms: + asset_name = base.action_manager.get_term(name).cfg.asset_name + if asset_name not in roots: + roots[asset_name] = base.arena_world.get_pose_w(asset_name)[env_id].detach().cpu().numpy() + geometry = {} + for key, image in images.items(): + term = getattr(camera_cfg, key) + sensor = base.scene[term.params["sensor_cfg"].name] + assert sensor.cfg.update_latest_camera_pose, "VLM calibration requires update_latest_camera_pose=True" + data = sensor.data + capture = ((data.info or {}).get("rgb") or {}).get("capture") + if capture: + from isaaclab.utils.math import convert_camera_frame_orientation_convention + + t_W_C = capture["pos_w"].torch[env_id].detach().cpu().numpy() + q_W_C = ( + convert_camera_frame_orientation_convention( + capture["quat_w_world"].torch[env_id : env_id + 1], origin="world", target="ros" + )[0] + .detach() + .cpu() + .numpy() + ) + intrinsics = capture["intrinsic_matrices"].torch[env_id] + else: + t_W_C = data.pos_w.torch[env_id].detach().cpu().numpy() + q_W_C = data.quat_w_ros.torch[env_id].detach().cpu().numpy() + intrinsics = data.intrinsic_matrices.torch[env_id] + # C is the optical frame (+X right, +Y down, +Z forward); R is an asset root. + poses = {} + for asset_name, T_W_R in roots.items(): + rotation_R_W = Rotation.from_quat(T_W_R[3:]).inv() + poses[asset_name] = { + "position_root_m": rotation_R_W.apply(t_W_C - T_W_R[:3]).tolist(), + "quaternion_root_xyzw": (rotation_R_W * Rotation.from_quat(q_W_C)).as_quat().tolist(), + } + geometry[key] = { + "image_size_px": [image.shape[1], image.shape[0]], + "intrinsic_matrix": intrinsics.detach().cpu().tolist(), + "optical_frame": "x_right_y_down_z_forward", + "position_world_m": t_W_C.tolist(), + "quaternion_world_xyzw": q_W_C.tolist(), + "poses_in_robot_roots": poses, + } + return geometry diff --git a/isaaclab_arena_vlm_agent/tests/conftest.py b/isaaclab_arena_vlm_agent/tests/conftest.py new file mode 100644 index 0000000000..e4774f89c2 --- /dev/null +++ b/isaaclab_arena_vlm_agent/tests/conftest.py @@ -0,0 +1,8 @@ +# Copyright (c) 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 + +"""Preserve pytest failures when Isaac Sim terminates the test process.""" + +from isaaclab_arena.tests.conftest import pytest_runtest_logreport, pytest_sessionstart # noqa: F401 diff --git a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py index 5ff9cdb089..47ff7c2a73 100644 --- a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py +++ b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py @@ -30,21 +30,38 @@ def inputs(monkeypatch): embodiment = SimpleNamespace( get_observation_cfg=lambda: SimpleNamespace( policy=SimpleNamespace(concatenate_terms=False, gripper_pos=SimpleNamespace(func=Mock())), - camera_obs=SimpleNamespace(front=SimpleNamespace(func=Mock(), params={"data_type": "rgb"})), + camera_obs=SimpleNamespace( + front=SimpleNamespace( + func=Mock(), params={"data_type": "rgb", "sensor_cfg": SimpleNamespace(name="sensor")} + ) + ), ), get_action_cfg=lambda: SimpleNamespace(arm=SimpleNamespace(asset_name="robot", body_name="tool")), ) world = SimpleNamespace(get_body_pose_in_root=Mock(return_value=torch.tensor([[0.4, 0, 0.3, 0, 0, 0, 1]] * 2))) + world.get_pose_w = Mock(return_value=torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]] * 2)) env = SimpleNamespace( unwrapped=SimpleNamespace( num_envs=2, - action_dim=8, device="cpu", + cfg=SimpleNamespace(observations=embodiment.get_observation_cfg()), + scene={ + "sensor": SimpleNamespace( + cfg=SimpleNamespace(update_latest_camera_pose=True), + data=SimpleNamespace( + pos_w=SimpleNamespace(torch=torch.zeros(2, 3)), + quat_w_ros=SimpleNamespace(torch=torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2)), + intrinsic_matrices=SimpleNamespace(torch=torch.eye(3).repeat(2, 1, 1)), + info={}, + ), + ) + }, arena_world=world, observation_manager=SimpleNamespace( active_terms={"policy": ["gripper_pos"]}, group_obs_concatenate={"policy": False} ), action_manager=SimpleNamespace( + total_action_dim=8, action=torch.zeros(2, 8), active_terms=["arm"], get_term=lambda name: SimpleNamespace(cfg=embodiment.get_action_cfg().arm), @@ -80,7 +97,7 @@ def _policy(mode, inputs, infer): policy = policy_type( config, inference_backend=SimpleNamespace(infer=infer), - observation_adapter=ImageStateObservationAdapter(embodiment), + observation_adapter=ImageStateObservationAdapter(), action_adapter=adapter, ) policy.set_task_description("Move the gripper") diff --git a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters_simulation.py b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters_simulation.py index 22a592659f..732461cabf 100644 --- a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters_simulation.py +++ b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters_simulation.py @@ -14,7 +14,7 @@ def _test_droid_absolute_pose_adapters(simulation_app): from isaaclab.utils.math import combine_frame_transforms - from isaaclab_arena.embodiments.droid.droid import DroidDifferentialIKEmbodiment + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteDifferentialIKEmbodiment from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment @@ -26,23 +26,17 @@ def _test_droid_absolute_pose_adapters(simulation_app): from isaaclab_arena_vlm_agent.pose_target_motion_generator import PoseTargetMotionGenerator # Use a translated, rotated root and separate environment origins to expose world/root confusion. - embodiment = DroidDifferentialIKEmbodiment( + embodiment = DroidAbsoluteDifferentialIKEmbodiment( initial_pose=Pose( position_xyz=(0.7, -0.4, 1.35), rotation_xyzw=(0.0, 0.0, math.sin(0.3), math.cos(0.3)), ) ) - def absolute_pose_actions(cfg): - cfg.actions.arm_action.controller.use_relative_mode = False - cfg.actions.arm_action.scale = 1.0 - return cfg - arena_env = IsaacLabArenaEnvironment( name="droid_vlm_adapter_test", scene=Scene(), embodiment=embodiment, - env_cfg_callback=absolute_pose_actions, ) env = ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg(num_envs=2, solve_relations=False)).make_registered() try: @@ -125,6 +119,15 @@ def _test_droid_joint_position_adapters(simulation_app): base = env.unwrapped adapter = AbsoluteChunkActionAdapter(1) assert adapter.is_robot_action_space_supported(env) + from isaaclab_arena_vlm_agent.droid_adapter import DroidVLMAdapter + + preset = DroidVLMAdapter() + for pose_adapter in ( + preset.make_step_action_adapter(), + preset.make_chunk_action_adapter(2), + preset.make_goal_action_adapter(0.01, 0.1), + ): + assert not pose_adapter.is_robot_action_space_supported(env) joint_names = tuple(f"panda_joint{i}" for i in range(1, 8)) targets = base.arena_world.get_joint_positions("robot", joint_names) actions = [] @@ -211,3 +214,22 @@ def term(cls, **kwargs): def test_absolute_controller_compatibility(): assert run_function_with_persistent_simulation_app(_test_absolute_controller_compatibility) + + +def _test_absolute_ik_config_is_consistent(simulation_app): + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteDifferentialIKEmbodiment + from isaaclab_arena.utils.physics_backend import PhysicsBackend + + for backend in PhysicsBackend: + embodiment = DroidAbsoluteDifferentialIKEmbodiment() + for configure in (False, True): + if configure: + embodiment.configure_physics_backend(backend) + assert embodiment.get_action_cfg() is embodiment.action_config + assert not embodiment.action_config.arm_action.controller.use_relative_mode + assert embodiment.action_config.arm_action.scale == 1.0 + return True + + +def test_absolute_ik_config_is_consistent(): + assert run_function_with_persistent_simulation_app(_test_absolute_ik_config_is_consistent) diff --git a/isaaclab_arena_vlm_agent/tests/test_image_utils.py b/isaaclab_arena_vlm_agent/tests/test_image_utils.py index 7185e4d726..5558a7c28e 100644 --- a/isaaclab_arena_vlm_agent/tests/test_image_utils.py +++ b/isaaclab_arena_vlm_agent/tests/test_image_utils.py @@ -21,9 +21,9 @@ def test_image_encoding_preserves_spatial_color_and_aspect_ratio(max_width_or_he pixels[:, :8, 0] = 255 pixels[:, 8:, 2] = 255 pixels[..., 3] = 255 - url = image_to_jpeg_data_url(pixels, max_width_or_height) + url, size = image_to_jpeg_data_url(pixels, max_width_or_height) image = Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))) - assert image.size == expected_size + assert image.size == size == expected_size assert image.mode == "RGB" left, right = image.getpixel((0, 0)), image.getpixel((image.width - 1, 0)) assert left[0] > left[2] + 100 diff --git a/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py b/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py index cbdb1641b2..ec12779d8b 100644 --- a/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py +++ b/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py @@ -5,6 +5,7 @@ """Verify embodiment-based state and image selection without robot-specific field names.""" +import numpy as np import torch from types import SimpleNamespace from unittest.mock import Mock @@ -27,8 +28,12 @@ def _inputs(): disabled_term=None, ) camera_cfg = SimpleNamespace( - head_color=SimpleNamespace(func=Mock(), params={"data_type": "rgb"}), - hand_color=SimpleNamespace(func=Mock(), params={"data_type": "rgb"}), + head_color=SimpleNamespace( + func=Mock(), params={"data_type": "rgb", "sensor_cfg": SimpleNamespace(name="camera_sensor")} + ), + hand_color=SimpleNamespace( + func=Mock(), params={"data_type": "rgb", "sensor_cfg": SimpleNamespace(name="camera_sensor")} + ), head_depth=SimpleNamespace(func=Mock(), params={"data_type": "distance_to_image_plane"}), ) action_cfg = SimpleNamespace( @@ -45,10 +50,25 @@ def get_pose(asset_name, body_name): x = 1.0 if asset_name == "left_robot" else 2.0 return torch.tensor([[x, 0, 0, 0, 0, 0, 1], [x + 10, 0, 0, 0, 0, 0, 1]]) - world = SimpleNamespace(get_body_pose_in_root=Mock(side_effect=get_pose)) + world = SimpleNamespace( + get_body_pose_in_root=Mock(side_effect=get_pose), + get_pose_w=Mock(return_value=torch.tensor([[0, 0, 0, 0, 0, 0, 1.0], [10, 0, 0, 0, 0, 2**-0.5, 2**-0.5]])), + ) + camera_data = SimpleNamespace( + pos_w=SimpleNamespace(torch=torch.tensor([[1.0, 2.0, 3.0], [10.0, 2.0, 3.0]])), + quat_w_ros=SimpleNamespace(torch=torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]])), + intrinsic_matrices=SimpleNamespace( + torch=torch.tensor([[[4.0, 0.0, 4.0], [0.0, 4.0, 2.0], [0.0, 0.0, 1.0]]] * 2) + ), + info={}, + ) env = SimpleNamespace( unwrapped=SimpleNamespace( num_envs=2, + cfg=SimpleNamespace(observations=embodiment.get_observation_cfg()), + scene={ + "camera_sensor": SimpleNamespace(cfg=SimpleNamespace(update_latest_camera_pose=True), data=camera_data) + }, arena_world=world, observation_manager=SimpleNamespace( active_terms={"policy": ["left_joints", "right_joints"]}, group_obs_concatenate={"policy": False} @@ -89,7 +109,7 @@ def test_selects_only_requested_modalities_and_configured_terms(adapter_type, in if not include_images: del observation["camera_obs"] embodiment.get_observation_cfg.return_value.camera_obs = None - adapter = adapter_type() if adapter_type is StateObservationAdapter else adapter_type(embodiment) + adapter = adapter_type() images, state = adapter.extract_vlm_input_from_observation(env, observation, 1) if include_state: assert state["policy"] == {"left_joints": [3.0, 4.0], "right_joints": [7.0, 8.0]} @@ -107,8 +127,8 @@ def test_selects_only_requested_modalities_and_configured_terms(adapter_type, in assert images["hand_color"].shape == (8, 4, 3) else: assert images == {} - if not include_state: - assert state == {} + if include_images: + assert set(state["camera_geometry"]) == set(images) def test_joint_only_state_does_not_require_command_body_queries(): @@ -129,12 +149,12 @@ def test_state_modality_requires_named_observation_terms(): def test_rgb_modality_requires_configured_rgb_views(): - embodiment, _, _ = _inputs() + embodiment, env, observation = _inputs() camera_cfg = embodiment.get_observation_cfg.return_value.camera_obs del camera_cfg.head_color del camera_cfg.hand_color with pytest.raises(AssertionError, match="configured RGB"): - ImageObservationAdapter(embodiment) + ImageObservationAdapter().extract_vlm_input_from_observation(env, observation, 0) def test_observation_base_requires_a_modality_implementation(): @@ -144,7 +164,7 @@ def test_observation_base_requires_a_modality_implementation(): def test_state_uses_live_managers_after_environment_configuration_changes(): embodiment, env, observation = _inputs() - adapter = ImageStateObservationAdapter(embodiment) + adapter = ImageStateObservationAdapter() env.unwrapped.observation_manager.active_terms["policy"] = ["task_only_term"] env.unwrapped.action_manager.active_terms = ["new_arm"] env.unwrapped.action_manager.get_term = lambda name: SimpleNamespace( @@ -154,3 +174,69 @@ def test_state_uses_live_managers_after_environment_configuration_changes(): assert state["policy"] == {"task_only_term": [1.0, 1.0, 1.0]} assert set(state["command_body_poses"]) == {"new_arm"} env.unwrapped.arena_world.get_body_pose_in_root.assert_called_once_with("new_robot", "new_tool") + + +def test_camera_optical_pose_uses_selected_environment_and_refreshes(): + _, env, observation = _inputs() + adapter = ImageObservationAdapter() + _, values = adapter.extract_vlm_input_from_observation(env, observation, 1) + camera = values["camera_geometry"]["head_color"] + assert camera["image_size_px"] == [8, 4] + assert camera["quaternion_world_xyzw"] == [0.0, 0.0, 1.0, 0.0] + pose = camera["poses_in_robot_roots"]["left_robot"] + np.testing.assert_allclose(pose["position_root_m"], [2, 0, 3], atol=1e-6) + np.testing.assert_allclose(pose["quaternion_root_xyzw"], [0, 0, 2**-0.5, 2**-0.5], atol=1e-6) + env.unwrapped.scene["camera_sensor"].data.pos_w.torch[1, 1] += 1 + _, values = adapter.extract_vlm_input_from_observation(env, observation, 1) + np.testing.assert_allclose( + values["camera_geometry"]["head_color"]["poses_in_robot_roots"]["left_robot"]["position_root_m"], + [3, 0, 3], + atol=1e-6, + ) + _, values = adapter.extract_vlm_input_from_observation(env, observation, 0) + assert values["camera_geometry"]["head_color"]["position_world_m"] == [1.0, 2.0, 3.0] + + +def test_delayed_camera_uses_capture_calibration(): + _, env, observation = _inputs() + data = env.unwrapped.scene["camera_sensor"].data + data.info = { + "rgb": { + "capture": { + "pos_w": SimpleNamespace(torch=torch.tensor([[4.0, 5.0, 6.0]] * 2)), + "quat_w_world": SimpleNamespace(torch=torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2)), + "intrinsic_matrices": SimpleNamespace(torch=torch.eye(3).repeat(2, 1, 1)), + } + } + } + _, values = ImageObservationAdapter().extract_vlm_input_from_observation(env, observation, 0) + camera = values["camera_geometry"]["head_color"] + assert camera["position_world_m"] == [4.0, 5.0, 6.0] + np.testing.assert_allclose(camera["intrinsic_matrix"], np.eye(3)) + from scipy.spatial.transform import Rotation + + # An identity world-convention camera looks along world +X, with image right along world -Y. + rotation_W_C = Rotation.from_quat(camera["quaternion_world_xyzw"]) + np.testing.assert_allclose(rotation_W_C.apply([0, 0, 1]), [1, 0, 0], atol=1e-6) + np.testing.assert_allclose(rotation_W_C.apply([1, 0, 0]), [0, -1, 0], atol=1e-6) + + +def test_camera_geometry_rejects_stale_pose_configuration(): + _, env, observation = _inputs() + env.unwrapped.scene["camera_sensor"].cfg.update_latest_camera_pose = False + with pytest.raises(AssertionError, match="update_latest_camera_pose=True"): + ImageObservationAdapter().extract_vlm_input_from_observation(env, observation, 0) + + +def test_camera_selection_follows_live_configuration_changes(): + _, env, observation = _inputs() + adapter = ImageObservationAdapter() + images, _ = adapter.extract_vlm_input_from_observation(env, observation, 0) + assert set(images) == {"head_color", "hand_color"} + cfg = env.unwrapped.cfg.observations.camera_obs + cfg.new_color = cfg.head_color + del cfg.head_color + del cfg.hand_color + observation["camera_obs"] = {"new_color": observation["camera_obs"]["head_color"]} + images, values = adapter.extract_vlm_input_from_observation(env, observation, 0) + assert set(images) == set(values["camera_geometry"]) == {"new_color"} diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py b/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py index 0be4a6147c..ce4876495e 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py @@ -68,7 +68,7 @@ def __init__(self, scale=1.0, action_dim=2): self.joint_names = tuple(f"joint_{i}" for i in range(action_dim)) def is_robot_action_space_supported(self, env): - return env.unwrapped.action_dim > 0 + return env.unwrapped.action_manager.total_action_dim > 0 def decode_command(self, response, arena_world, env_id): positions = arena_world.get_joint_positions("robot", self.joint_names)[env_id].cpu().numpy() @@ -116,7 +116,9 @@ def make_policy(backend, **kwargs): def inputs(action_dim=2): - env = SimpleNamespace(unwrapped=SimpleNamespace(num_envs=2, device="cpu", action_dim=action_dim)) + env = SimpleNamespace( + unwrapped=SimpleNamespace(num_envs=2, device="cpu", action_manager=SimpleNamespace(total_action_dim=action_dim)) + ) positions = torch.arange(2 * action_dim, dtype=torch.float32).reshape(2, action_dim) def get_joint_positions(scene_key, joint_names): @@ -131,6 +133,28 @@ def get_joint_positions(scene_key, joint_names): return env, {"position": positions, CAMERA_OBSERVATION_GROUP: {"front_view": images}} +@pytest.mark.parametrize("max_size,expected_size", [(8, (8, 4)), (32, (17, 9))]) +def test_request_intrinsics_match_encoded_pixels(backend, max_size, expected_size): + policy = make_policy(backend) + policy.config.image_max_width_or_height = max_size + env, observation = inputs() + pixels = np.zeros((9, 17, 3), dtype=np.uint8) + intrinsics = np.array([[12.0, 0.0, 8.0], [0.0, 10.0, 4.0], [0.0, 0.0, 1.0]]) + metadata = {"camera_geometry": {"front_view": {"image_size_px": [17, 9], "intrinsic_matrix": intrinsics.tolist()}}} + original_metadata = copy.deepcopy(metadata) + policy.observation_adapter.extract_vlm_input_from_observation = lambda *args: ({"front_view": pixels}, metadata) + content = policy._request_content(env, observation, 0) + assert policy._request_content(env, observation, 0) == content + assert metadata == original_metadata + geometry = json.loads(content[1]["text"])["camera_geometry"]["front_view"] + encoded = Image.open(io.BytesIO(base64.b64decode(content[3]["image_url"]["url"].split(",", 1)[1]))) + assert encoded.size == expected_size == tuple(geometry["image_size_px"]) + p_C = np.array([0.2, -0.1, 2.0]) + original_pixel = (intrinsics @ p_C)[:2] / p_C[2] + resized_pixel = (np.array(geometry["intrinsic_matrix"]) @ p_C)[:2] / p_C[2] + np.testing.assert_allclose(resized_pixel, original_pixel * np.array(expected_size) / [17, 9]) + + @pytest.mark.parametrize("action_dim", [2, 5]) def test_policy_uses_adapter_actions_for_each_environment(backend, action_dim): client = backend.client @@ -175,7 +199,7 @@ def create(**kwargs): assert all(len(call["messages"]) == 2 for call in calls) finally: policy.close() - client.close.assert_not_called() # The backend belongs to the caller, not the policy. + client.close.assert_called_once() @pytest.mark.parametrize( @@ -234,7 +258,7 @@ def test_environment_validation_precedes_inference(backend): client = backend.client policy = make_policy(backend) env, observation = inputs() - env.unwrapped.action_dim = 0 + env.unwrapped.action_manager.total_action_dim = 0 try: with pytest.raises(AssertionError, match="Robot action space"): policy.get_action(env, observation) @@ -265,7 +289,7 @@ def test_goal_config_rejects_invalid_step_budget(max_steps): class JointChunkAdapter(ChunkActionAdapter): def is_robot_action_space_supported(self, env): - return env.unwrapped.action_dim > 0 + return env.unwrapped.action_manager.total_action_dim > 0 def __init__(self, chunk_size: int, action_dim=2): super().__init__(chunk_size) @@ -293,7 +317,7 @@ def __init__(self, action_dim=2, *, motion_generator=None): self.joint_names = tuple(f"joint_{i}" for i in range(action_dim)) def is_robot_action_space_supported(self, env): - return env.unwrapped.action_dim > 0 + return env.unwrapped.action_manager.total_action_dim > 0 def decode_command(self, response, arena_world, env_id): positions = arena_world.get_joint_positions("robot", self.joint_names)[env_id].cpu().numpy() diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py new file mode 100644 index 0000000000..7e027ee2b3 --- /dev/null +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py @@ -0,0 +1,294 @@ +# Copyright (c) 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 + +"""Check registered policy construction, execution strategies, and client ownership.""" + +import torch +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from isaaclab_arena_vlm_agent import droid_adapter, vlm_agent_policy +from isaaclab_arena_vlm_agent.absolute_action_adapter import AbsoluteChunkActionAdapter, AbsoluteStepActionAdapter +from isaaclab_arena_vlm_agent.command_indices import AbsoluteCommandIndices, GripperCommandIndex, PoseCommandIndices +from isaaclab_arena_vlm_agent.vlm_agent_policy import ( + VLMChunkPolicy, + VLMChunkPolicyCfg, + VLMGoalPolicy, + VLMGoalPolicyCfg, + VLMStepPolicy, + VLMStepPolicyCfg, +) + +POLICIES = [ + (VLMStepPolicy, VLMStepPolicyCfg), + (VLMChunkPolicy, VLMChunkPolicyCfg), + (VLMGoalPolicy, VLMGoalPolicyCfg), +] + + +@pytest.fixture(autouse=True) +def gripper_measurements(monkeypatch): + gripper = SimpleNamespace( + get_closing_error_m=lambda env: torch.zeros(2), + get_opening_width_m=lambda world: torch.zeros(2), + ) + monkeypatch.setattr(droid_adapter, "_make_gripper", lambda: gripper) + + +def _config(cfg_type, modality="image_state"): + options = {"max_steps_per_goal": 2} if cfg_type is VLMGoalPolicyCfg else {} + if cfg_type is VLMChunkPolicyCfg: + options["chunk_size"] = 2 + cfg = cfg_type(system_prompt="Return absolute targets.", model="test-model", **options) + cfg.observation_modality = modality + return cfg + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +@pytest.mark.parametrize("modality", ["image", "state", "image_state"]) +def test_configured_execution_and_partial_reset(monkeypatch, policy_type, cfg_type, modality): + backend = Mock() + replies = [] + for x in (1.0, 3.0, 5.0, 7.0): + command = {"values": [x, 0, 0, 0, 0, 0, 1, 0]} + if policy_type is VLMChunkPolicy: + command = {"commands": [command, {"values": [x + 1, 0, 0, 0, 0, 0, 1, 0]}]} + replies.append(command) + backend.infer.side_effect = replies + constructor = Mock(return_value=backend) + monkeypatch.setattr(vlm_agent_policy, "InferenceBackend", constructor) + indices = AbsoluteCommandIndices( + 8, + (PoseCommandIndices(slice(0, 7), "robot", "base_link"),), + (), + (GripperCommandIndex(7, "gripper_action", 0.0, 1.0),), + ) + for module in ("absolute_action_adapter", "pose_target_motion_generator"): + monkeypatch.setattr(f"isaaclab_arena_vlm_agent.{module}.read_absolute_command_indices", lambda env: indices) + observation_cfg = SimpleNamespace( + camera_obs=SimpleNamespace( + live_camera=SimpleNamespace( + func=Mock(), params={"data_type": "rgb", "sensor_cfg": SimpleNamespace(name="sensor")} + ) + ) + ) + env = SimpleNamespace( + unwrapped=SimpleNamespace( + num_envs=2, + device="cpu", + scene={ + "sensor": SimpleNamespace( + cfg=SimpleNamespace(update_latest_camera_pose=True), + data=SimpleNamespace( + pos_w=SimpleNamespace(torch=torch.zeros(2, 3)), + quat_w_ros=SimpleNamespace(torch=torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2)), + intrinsic_matrices=SimpleNamespace(torch=torch.eye(3).repeat(2, 1, 1)), + info={}, + ), + ) + }, + arena_world=SimpleNamespace( + get_body_pose_in_root=lambda *args: torch.tensor([[0, 0, 0, 0, 0, 0, 1.0]] * 2) + ), + cfg=SimpleNamespace(observations=observation_cfg), + observation_manager=SimpleNamespace( + active_terms={"policy": ["joint_pos"]}, group_obs_concatenate={"policy": False} + ), + action_manager=SimpleNamespace(active_terms=[], total_action_dim=8, action=torch.zeros(2, 8)), + episode_length_buf=torch.zeros(2, dtype=torch.long), + ) + ) + observations = { + "policy": {"joint_pos": torch.zeros(2, 1)}, + "camera_obs": {"live_camera": torch.zeros(2, 4, 4, 3, dtype=torch.uint8)}, + } + policy = policy_type(_config(cfg_type, modality)) + policy.reset() + policy.set_task_description("Move the robot") + torch.testing.assert_close(policy.get_action(env, observations)[:, 0], torch.tensor([1.0, 3.0])) + content = backend.infer.call_args.args[0].messages[1]["content"] + assert any(item["type"] == "image_url" for item in content) == (modality != "state") + assert any('"joint_pos"' in item.get("text", "") for item in content) == (modality != "image") + assert backend.infer.call_args.args[0].messages[0]["content"] == "Return absolute targets." + policy.reset(torch.tensor([0])) + expected_second = {VLMStepPolicy: 7.0, VLMChunkPolicy: 4.0, VLMGoalPolicy: 3.0}[policy_type] + torch.testing.assert_close(policy.get_action(env, observations)[:, 0], torch.tensor([5.0, expected_second])) + assert backend.infer.call_count == (4 if policy_type is VLMStepPolicy else 3) + if policy_type is VLMGoalPolicy: + # Only env 1 has exhausted the two-step goal budget after the partial reset. + torch.testing.assert_close(policy.get_action(env, observations)[:, 0], torch.tensor([5.0, 7.0])) + assert backend.infer.call_count == 4 + assert constructor.call_args.kwargs["model"] == "test-model" + policy.close() + backend.close.assert_called_once() + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +def test_policy_closes_injected_backend(monkeypatch, policy_type, cfg_type): + constructor = Mock() + monkeypatch.setattr(vlm_agent_policy, "InferenceBackend", constructor) + backend = Mock() + policy = policy_type(_config(cfg_type), inference_backend=backend) + policy.close() + constructor.assert_not_called() + backend.close.assert_called_once() + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +def test_invalid_adapter_rejected_before_client_creation(monkeypatch, policy_type, cfg_type): + constructor = Mock() + monkeypatch.setattr(vlm_agent_policy, "InferenceBackend", constructor) + cfg = _config(cfg_type) + adapter = AbsoluteChunkActionAdapter(1) if policy_type is not VLMChunkPolicy else AbsoluteStepActionAdapter() + with pytest.raises(AssertionError): + policy_type(cfg, action_adapter=adapter) + constructor.assert_not_called() + + +def test_invalid_schema_rejected_before_client_creation(monkeypatch): + constructor = Mock() + monkeypatch.setattr(vlm_agent_policy, "InferenceBackend", constructor) + monkeypatch.setattr(vlm_agent_policy, "build_strict_schema", Mock(side_effect=RuntimeError("invalid schema"))) + with pytest.raises(RuntimeError, match="invalid schema"): + VLMStepPolicy(_config(VLMStepPolicyCfg)) + constructor.assert_not_called() + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +def test_example_yaml_policy_selectors_and_round_trip(monkeypatch, tmp_path, policy_type, cfg_type): + from isaaclab_arena.assets import registries + from isaaclab_arena.evaluation import arena_experiment_config_loader as loader + from isaaclab_arena.hydra.typed_experiment_serializer import serialize_arena_experiment_to_yaml + from isaaclab_arena_environments.pick_and_place_maple_table_environment import PickAndPlaceMapleTableEnvironmentCfg + + monkeypatch.setattr(registries, "ensure_assets_registered", lambda: None) + monkeypatch.setattr( + loader, + "_registered_environment_cfg_types", + lambda: {"pick_and_place_maple_table": PickAndPlaceMapleTableEnvironmentCfg}, + ) + class_path = f"isaaclab_arena_vlm_agent.vlm_agent_policy.{policy_type.__name__}" + for selector in (policy_type.name, class_path): + assert loader._resolve_policy_cfg_type_from_name_or_class_path(selector) is cfg_type + root = Path(__file__).resolve().parents[2] + run_name = f"droid_pnp_{policy_type.name}" + overrides = [] + if policy_type is VLMGoalPolicy: + overrides = [f"runs.{run_name}.policy.max_steps_per_goal=5"] + elif policy_type is VLMChunkPolicy: + overrides = [f"runs.{run_name}.policy.chunk_size=2"] + experiment = loader.load_arena_experiment_from_config_file( + root / f"isaaclab_arena_environments/experiment_configs/{run_name}_experiment.yaml", + device="cpu", + overrides=overrides, + ) + assert set(experiment.runs) == {run_name} + run = experiment.runs[run_name] + assert isinstance(run.policy, cfg_type) + assert run.policy.model == "openai/openai/gpt-6-astra" + assert run.environment_builder.device == "cpu" + assert registries.PolicyRegistry().get_policy_type_for_cfg(run.policy) is policy_type + if policy_type is VLMGoalPolicy: + assert run.policy.max_steps_per_goal == 5 + elif policy_type is VLMChunkPolicy: + assert run.policy.chunk_size == 2 + round_trip_path = tmp_path / "experiment.yaml" + round_trip_path.write_text(serialize_arena_experiment_to_yaml(experiment)) + restored = loader.load_arena_experiment_from_config_file(round_trip_path, device="cpu") + assert restored == experiment + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +def test_preset_supplies_prompt_and_strategy(monkeypatch, policy_type, cfg_type): + monkeypatch.setattr(vlm_agent_policy, "InferenceBackend", Mock()) + cfg = _config(cfg_type) + cfg.system_prompt = None + policy = policy_type(cfg) + assert "base_link" in policy._system_prompt + if policy_type is VLMChunkPolicy: + assert "exactly 2 commands" in policy._system_prompt + assert policy.action_adapter.chunk_size == 2 + else: + assert "Return one command" in policy._system_prompt + if policy_type is VLMGoalPolicy: + assert "gripper_action" in policy.action_adapter.motion_generator.grippers + policy.close() + + +@pytest.mark.parametrize( + "options", [{"vlm_embodiment_adapter": "unknown"}, {"observation_modality": "unknown"}, {"chunk_size": 0}] +) +def test_invalid_flat_settings(options): + with pytest.raises((AssertionError, ValueError)): + VLMChunkPolicyCfg(**{"chunk_size": 2, **options}) + + +@pytest.mark.parametrize("class_name", ["VLMStepPolicy", "VLMChunkPolicy", "VLMGoalPolicy"]) +def test_policy_class_path_loads_in_fresh_process(class_name): + import subprocess + import sys + + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; " + "from isaaclab_arena.evaluation.arena_experiment_config_loader " + "import _resolve_policy_cfg_type_from_name_or_class_path as resolve; " + "assert 'isaaclab_arena_vlm_agent.vlm_agent_policy' not in sys.modules; " + f"cfg = resolve('isaaclab_arena_vlm_agent.vlm_agent_policy.{class_name}'); " + "assert cfg.__module__ == 'isaaclab_arena_vlm_agent.vlm_agent_policy'" + ), + ], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_chunk_size_is_required(): + with pytest.raises(TypeError, match="chunk_size"): + VLMChunkPolicyCfg() + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +@pytest.mark.parametrize("layout", ["joints", "wrong_body", "wrong_asset", "gripper_first", "signed_gripper"]) +def test_droid_preset_rejects_incompatible_controller_before_inference(monkeypatch, policy_type, cfg_type, layout): + from dataclasses import replace + + indices = AbsoluteCommandIndices( + 8, + (PoseCommandIndices(slice(0, 7), "robot", "base_link"),), + (), + (GripperCommandIndex(7, "gripper_action", 0.0, 1.0),), + ) + if layout == "joints": + indices = replace(indices, poses=(), joints=(slice(0, 7),)) + elif layout == "wrong_body": + indices = replace(indices, poses=(replace(indices.poses[0], body_name="tool"),)) + elif layout == "wrong_asset": + indices = replace(indices, poses=(replace(indices.poses[0], asset_name="other_robot"),)) + elif layout == "gripper_first": + indices = replace( + indices, + poses=(replace(indices.poses[0], indices=slice(1, 8)),), + grippers=(replace(indices.grippers[0], index=0),), + ) + else: + indices = replace(indices, grippers=(replace(indices.grippers[0], open_value=1, close_value=-1),)) + monkeypatch.setattr( + "isaaclab_arena_vlm_agent.absolute_action_adapter.read_absolute_command_indices", lambda env: indices + ) + backend = Mock() + policy = policy_type(_config(cfg_type), inference_backend=backend) + policy.set_task_description("Move DROID") + with pytest.raises(AssertionError, match="unsupported"): + policy.get_action(SimpleNamespace(), {}) + backend.infer.assert_not_called() diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py new file mode 100644 index 0000000000..6156275723 --- /dev/null +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py @@ -0,0 +1,119 @@ +# Copyright (c) 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 + +"""Verify goal-policy execution and moving-camera calibration in an isolated simulator.""" + +import sys + +import pytest + +from isaaclab_arena.tests.utils.subprocess import run_subprocess + + +def _test_vlm_experiment(simulation_app): + import base64 + import io + import json + import numpy as np + import tempfile + from pathlib import Path + from scipy.spatial.transform import Rotation + from unittest.mock import Mock, patch + + from PIL import Image + + from isaaclab_arena.assets.registries import PolicyRegistry + from isaaclab_arena.embodiments.droid.droid import DroidCameraCfg + from isaaclab_arena.evaluation.arena_experiment_config_loader import load_arena_experiment_from_config_file + from isaaclab_arena.evaluation.arena_run import RunStatus + from isaaclab_arena.evaluation.run_execution import build_and_run + from isaaclab_arena.tests.utils.constants import TestConstants + + config_path = Path(TestConstants.arena_environments_dir) / "experiment_configs/droid_pnp_vlm_goal_experiment.yaml" + run_name = "droid_pnp_vlm_goal" + experiment = load_arena_experiment_from_config_file( + config_path, + device="cuda:0", + overrides=[ + f"runs.{run_name}.rollout_limit.num_episodes=null", + f"runs.{run_name}.rollout_limit.num_steps=3", + f"runs.{run_name}.policy.model=test-model", + f"runs.{run_name}.policy.max_steps_per_goal=1", + ], + ) + run = experiment.runs[run_name] + assert PolicyRegistry().get_policy_type_for_cfg(run.policy) is PolicyRegistry().get_policy("vlm_goal") + backend = Mock() + wrist_positions = [] + wrist_offset = DroidCameraCfg().wrist_camera.offset + + def move_nearby(request): + content = request.messages[1]["content"] + assert any(item["type"] == "image_url" for item in content) + state = json.loads(content[1]["text"]) + assert "All quaternions in input and output use XYZW" in request.messages[0]["content"] + assert len(state["camera_geometry"]) == 3 + for index in range(2, len(content), 2): + name = content[index]["text"].removeprefix("camera=") + geometry = state["camera_geometry"][name] + url = content[index + 1]["image_url"]["url"] + image = Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))) + assert tuple(geometry["image_size_px"]) == image.size == (384, 216) + assert geometry["optical_frame"] == "x_right_y_down_z_forward" + K = np.array(geometry["intrinsic_matrix"]) + assert K[0, 0] > 0 and K[1, 1] > 0 + np.testing.assert_allclose(K[2], [0, 0, 1]) + np.testing.assert_allclose(np.linalg.norm(geometry["quaternion_world_xyzw"]), 1, atol=1e-5) + # DROID's fixed root is identity in this example, so world and root camera poses agree. + root = geometry["poses_in_robot_roots"]["robot"] + np.testing.assert_allclose(root["position_root_m"], geometry["position_world_m"], atol=1e-5) + np.testing.assert_allclose( + Rotation.from_quat(root["quaternion_root_xyzw"]).as_matrix(), + Rotation.from_quat(geometry["quaternion_world_xyzw"]).as_matrix(), + atol=1e-5, + ) + pose = state["command_body_poses"]["arm_action"] + np.testing.assert_allclose(state["policy"]["eef_quat"], pose["quaternion_root_xyzw"], atol=1e-5) + # Reset can return an image captured before the joint reset. After a rendered step, + # check the moving camera against its measured parent link, not its own cached data. + wrist = state["camera_geometry"]["wrist_camera_rgb"] + if wrist_positions: + rotation_W_E = Rotation.from_quat(state["policy"]["eef_quat"]) + t_W_C = rotation_W_E.apply(wrist_offset.pos) + state["policy"]["eef_pos"] + np.testing.assert_allclose(wrist["position_world_m"], t_W_C, atol=1e-4) + rotation_W_C = rotation_W_E * Rotation.from_quat(wrist_offset.rot) * Rotation.from_euler("x", np.pi) + # Allow sub-milliradian rounding in the authored USD mount quaternion. + rotation_error = rotation_W_C.inv() * Rotation.from_quat(wrist["quaternion_world_xyzw"]) + assert rotation_error.magnitude() < 5e-4 + wrist_positions.append(wrist["position_world_m"]) + command = {"values": pose["position_root_m"] + pose["quaternion_root_xyzw"] + [0.0]} + command["values"][0] += 0.01 + return command + + backend.infer.side_effect = move_nearby + with tempfile.TemporaryDirectory() as output_dir: + with patch("isaaclab_arena_vlm_agent.vlm_agent_policy.InferenceBackend", return_value=backend): + result = build_and_run(run, output_dir) + assert result.status is RunStatus.COMPLETED + assert backend.infer.call_count == 3 + assert np.linalg.norm(np.array(wrist_positions[-1]) - wrist_positions[0]) > 1e-4 + backend.close.assert_called_once() + return True + + +@pytest.mark.with_cameras +@pytest.mark.with_subprocess +def test_vlm_experiment(): + # Other VLM tests initialize the parent SimulationApp without camera support. + run_subprocess([ + sys.executable, + "-c", + ( + "from isaaclab_arena.tests.utils.persistent_simulation_app import" + " run_function_with_persistent_simulation_app; from" + " isaaclab_arena_vlm_agent.tests.test_vlm_policy_configuration_simulation import _test_vlm_experiment;" + " assert run_function_with_persistent_simulation_app(_test_vlm_experiment, enable_cameras=True)" + ), + ]) diff --git a/isaaclab_arena_vlm_agent/vlm_agent_policy.py b/isaaclab_arena_vlm_agent/vlm_agent_policy.py index 9671ed7b58..ce64f428a5 100644 --- a/isaaclab_arena_vlm_agent/vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/vlm_agent_policy.py @@ -13,28 +13,53 @@ import torch from abc import abstractmethod from collections import deque +from copy import deepcopy from dataclasses import dataclass, field +from enum import Enum from typing import TYPE_CHECKING, Any from pydantic import ValidationError -from isaaclab_arena.inference.backend import InferenceBackend, InferenceRequest, build_strict_schema +from isaaclab_arena.assets.register import register_policy +from isaaclab_arena.inference.backend import MAX_RETRIES_LIMIT, InferenceBackend, InferenceRequest, build_strict_schema from isaaclab_arena.policy.policy_base import PolicyBase, PolicyCfg from isaaclab_arena_vlm_agent.action_adapter import AgentActionAdapter, ChunkActionAdapter, GoalActionAdapter -from isaaclab_arena_vlm_agent.image_utils import image_to_jpeg_data_url -from isaaclab_arena_vlm_agent.observation_adapter import VLMObservationAdapter +from isaaclab_arena_vlm_agent.image_utils import image_to_jpeg_data_url, resize_camera_geometry +from isaaclab_arena_vlm_agent.observation_adapter import ( + ImageObservationAdapter, + ImageStateObservationAdapter, + StateObservationAdapter, + VLMObservationAdapter, +) from isaaclab_arena_vlm_agent.response_errors import InvalidAgentResponseError if TYPE_CHECKING: from isaaclab_arena.environments.arena_world import ArenaWorld +class ObservationModality(str, Enum): + """Observations included in each VLM request.""" + + # Hydra accepts enum member names, so retain the lowercase YAML vocabulary. + image = "image" + state = "state" + image_state = "image_state" + + def get_adapter_type(self) -> type[VLMObservationAdapter]: + """Return the observation adapter selected by this modality.""" + return { + self.image: ImageObservationAdapter, + self.state: StateObservationAdapter, + self.image_state: ImageStateObservationAdapter, + }[self] + + @dataclass class VLMPolicyCfg(PolicyCfg): """Configure policy prompts, image encoding, and action-validation attempts.""" - system_prompt: str - """Required system message defining the task, units, and response contract.""" + system_prompt: str | None + """System message override, or None to use the selected embodiment adapter's prompt.""" image_max_width_or_height: int = 384 """Maximum image width or height sent to the endpoint, in pixels. @@ -46,21 +71,91 @@ class VLMPolicyCfg(PolicyCfg): max_attempts: int = 3 """Maximum response-validation attempts per environment step.""" + vlm_embodiment_adapter: str | None = None + """Named action and prompt preset; droid supports the droid_abs_ik embodiment.""" + + observation_modality: ObservationModality = ObservationModality.image_state + """One of image, state, or image_state; ignored when an observation adapter is injected.""" + + endpoint: str | None = None + """Inference endpoint preset, or None to use ARENA_INFERENCE_ENDPOINT.""" + + model: str | None = None + """Optional model override; image observations require a vision-capable model.""" + + base_url: str | None = None + """Optional URL override; credentials come from the endpoint's environment variable.""" + + temperature: float = 0.2 + """Sampling temperature when supported by the endpoint.""" + + max_tokens: int = 4096 + """Maximum output tokens per inference request.""" + + max_retries: int = 3 + """Additional transport attempts, separate from response-validation max_attempts.""" + def __post_init__(self): - assert ( - isinstance(self.system_prompt, str) and self.system_prompt.strip() - ), "Set policy.system_prompt in the experiment" + if self.system_prompt is None: + assert self.vlm_embodiment_adapter is not None, "Set policy.system_prompt or policy.vlm_embodiment_adapter" + else: + assert isinstance(self.system_prompt, str) and self.system_prompt.strip(), "Set policy.system_prompt" + assert self.vlm_embodiment_adapter in (None, "droid"), "Unknown vlm_embodiment_adapter; expected 'droid'" + self.observation_modality = ObservationModality(self.observation_modality) assert self.image_max_width_or_height > 0 assert isinstance(self.max_attempts, int) and self.max_attempts >= 1 + assert type(self.max_tokens) is int and self.max_tokens > 0 + assert type(self.max_retries) is int and 0 <= self.max_retries < MAX_RETRIES_LIMIT + + +@dataclass +class VLMStepPolicyCfg(VLMPolicyCfg): + """Configure one command per prediction.""" + + system_prompt: str | None = None + """Optional override for the selected embodiment's single-command prompt.""" + + vlm_embodiment_adapter: str | None = "droid" + """Named action and prompt preset.""" + + +@dataclass +class VLMChunkPolicyCfg(VLMPolicyCfg): + """Configure command chunks executed between predictions.""" + + system_prompt: str | None = None + """Optional override for the selected embodiment's chunk prompt.""" + + vlm_embodiment_adapter: str | None = "droid" + """Named action and prompt preset.""" + + chunk_size: int = field(kw_only=True) + """Number of commands per prediction; also used by the default prompt.""" + + def __post_init__(self): + super().__post_init__() + assert type(self.chunk_size) is int and self.chunk_size > 0, "chunk_size must be a positive integer" @dataclass class VLMGoalPolicyCfg(VLMPolicyCfg): """Configure VLM goal prediction and the execution budget for each goal.""" + system_prompt: str | None = None + """Optional override for the selected embodiment's goal prompt.""" + + vlm_embodiment_adapter: str | None = "droid" + """Named action, motion-generator, gripper, and prompt preset.""" + max_steps_per_goal: int = field(kw_only=True) """Maximum commands to execute before requesting a new goal.""" + position_tolerance_m: float = 0.01 + """Maximum translation error for the configured pose motion generator.""" + + rotation_tolerance_rad: float = 0.1 + """Maximum rotation error for the configured pose motion generator.""" + def __post_init__(self): super().__post_init__() assert ( @@ -78,27 +173,50 @@ def __init__( self, config: VLMPolicyCfg, *, - inference_backend: InferenceBackend, - observation_adapter: VLMObservationAdapter, + inference_backend: InferenceBackend | None = None, + observation_adapter: VLMObservationAdapter | None = None, action_adapter: AgentActionAdapter, ): - """Create a policy with a shared inference backend and stateless adapters. + """Create configured dependencies, or use explicitly injected dependencies. Args: config: Prompt, image encoding, and action-validation settings. - inference_backend: Caller-owned backend shared with other inference steps. - The caller is responsible for closing its client. - observation_adapter: Adapter selecting camera images and prompt values for the VLM. + inference_backend: Backend whose ownership transfers to this policy, or None to create one. + observation_adapter: Adapter selecting images and state, or None to construct it from config. action_adapter: Adapter validating responses and converting commands to native actions. """ super().__init__(config) + if observation_adapter is None: + adapter_type = ObservationModality(config.observation_modality).get_adapter_type() + observation_adapter = adapter_type() assert isinstance(observation_adapter, VLMObservationAdapter) assert isinstance(action_adapter, AgentActionAdapter) self.observation_adapter = observation_adapter self.action_adapter = action_adapter self._response_schema = build_strict_schema(action_adapter.response_model) - self._inference_backend = inference_backend + self._system_prompt = config.system_prompt + if self._system_prompt is None: + self._system_prompt = _get_vlm_embodiment_adapter(config.vlm_embodiment_adapter).get_system_prompt( + action_adapter + ) self.task_description = None + # Construct the client last so invalid adapters or schemas cannot leak a connection. + self._inference_backend = ( + inference_backend + if inference_backend is not None + else InferenceBackend( + endpoint=config.endpoint, + model=config.model, + base_url=config.base_url, + temperature=config.temperature, + max_tokens=config.max_tokens, + max_retries=config.max_retries, + ) + ) + + def close(self) -> None: + """Close the inference backend owned by this policy.""" + self._inference_backend.close() def get_action(self, env: gym.Env, observation: dict[str, Any]) -> torch.Tensor: """Return a batch of native actions using the policy execution strategy. @@ -131,15 +249,25 @@ def _get_command(self, env: gym.Env, observation: dict[str, Any], env_id: int) - def _request_content(self, env: gym.Env, observation: dict[str, Any], env_id: int) -> list[dict[str, Any]]: """Encode selected images and assemble the task and observation values into chat content.""" images, prompt_values = self.observation_adapter.extract_vlm_input_from_observation(env, observation, env_id) + # Injected adapters may reuse calibration dictionaries across predictions. + prompt_values = deepcopy(prompt_values) + encoded_images = {} + for camera_name, image in images.items(): + url, size = image_to_jpeg_data_url(image, self.config.image_max_width_or_height) + encoded_images[camera_name] = url + if camera_name in prompt_values.get("camera_geometry", {}): + resize_camera_geometry( + prompt_values["camera_geometry"][camera_name], (image.shape[1], image.shape[0]), size + ) content = [{"type": "text", "text": json.dumps({"task": self.task_description})}] if prompt_values: content.append({"type": "text", "text": json.dumps(prompt_values, allow_nan=False)}) - for camera_name, image in images.items(): + for camera_name, url in encoded_images.items(): content.extend([ {"type": "text", "text": f"camera={camera_name}"}, { "type": "image_url", - "image_url": {"url": image_to_jpeg_data_url(image, self.config.image_max_width_or_height)}, + "image_url": {"url": url}, }, ]) return content @@ -147,7 +275,7 @@ def _request_content(self, env: gym.Env, observation: dict[str, Any], env_id: in def _infer(self, content: list[dict], arena_world: ArenaWorld, env_id: int) -> np.ndarray: """Request a command, allowing the model to correct locally rejected responses.""" messages = [ - {"role": "system", "content": self.config.system_prompt}, + {"role": "system", "content": self._system_prompt}, {"role": "user", "content": content}, ] rejection: InvalidAgentResponseError | None = None @@ -185,17 +313,22 @@ def _infer(self, content: list[dict], arena_world: ArenaWorld, env_id: int) -> n raise rejection -class VLMStepPolicy(VLMPolicy): +@register_policy +class VLMStepPolicy(VLMPolicy, PolicyBase[VLMStepPolicyCfg]): """Predict one action per environment from fresh observations on every step.""" + name = "vlm_step" + def __init__( self, config: VLMPolicyCfg, *, - inference_backend: InferenceBackend, - observation_adapter: VLMObservationAdapter, - action_adapter: AgentActionAdapter, + inference_backend: InferenceBackend | None = None, + observation_adapter: VLMObservationAdapter | None = None, + action_adapter: AgentActionAdapter | None = None, ): + if action_adapter is None: + action_adapter = _get_vlm_embodiment_adapter(config.vlm_embodiment_adapter).make_step_action_adapter() assert not isinstance(action_adapter, (ChunkActionAdapter, GoalActionAdapter)), "Expected a single-step adapter" super().__init__( config, @@ -208,17 +341,25 @@ def _get_command(self, env: gym.Env, observation: dict[str, Any], env_id: int) - return self._infer(self._request_content(env, observation, env_id), env.unwrapped.arena_world, env_id) -class VLMChunkPolicy(VLMPolicy): +@register_policy +class VLMChunkPolicy(VLMPolicy, PolicyBase[VLMChunkPolicyCfg]): """Execute a validated command chunk before requesting another from current observations.""" + name = "vlm_chunk" + def __init__( self, config: VLMPolicyCfg, *, - inference_backend: InferenceBackend, - observation_adapter: VLMObservationAdapter, - action_adapter: ChunkActionAdapter, + inference_backend: InferenceBackend | None = None, + observation_adapter: VLMObservationAdapter | None = None, + action_adapter: ChunkActionAdapter | None = None, ): + if action_adapter is None: + assert isinstance(config, VLMChunkPolicyCfg), "Configure VLMChunkPolicyCfg or inject an action adapter" + action_adapter = _get_vlm_embodiment_adapter(config.vlm_embodiment_adapter).make_chunk_action_adapter( + config.chunk_size + ) assert isinstance(action_adapter, ChunkActionAdapter) super().__init__( config, @@ -243,17 +384,24 @@ def reset(self, env_ids: torch.Tensor | None = None) -> None: self._chunks.pop(env_id, None) +@register_policy class VLMGoalPolicy(VLMPolicy, PolicyBase[VLMGoalPolicyCfg]): """Use a motion generator to execute decoded goals until completion or a step limit.""" + name = "vlm_goal" + def __init__( self, config: VLMGoalPolicyCfg, *, - inference_backend: InferenceBackend, - observation_adapter: VLMObservationAdapter, - action_adapter: GoalActionAdapter, + inference_backend: InferenceBackend | None = None, + observation_adapter: VLMObservationAdapter | None = None, + action_adapter: GoalActionAdapter | None = None, ): + if action_adapter is None: + action_adapter = _get_vlm_embodiment_adapter(config.vlm_embodiment_adapter).make_goal_action_adapter( + config.position_tolerance_m, config.rotation_tolerance_rad + ) assert isinstance(action_adapter, GoalActionAdapter) super().__init__( config, @@ -278,7 +426,7 @@ def _get_command(self, env: gym.Env, observation: dict[str, Any], env_id: int) - self.action_adapter.motion_generator.compute_command(goal, env, env_id), dtype=np.float32, ) - expected_dim = env.unwrapped.action_dim + expected_dim = env.unwrapped.action_manager.total_action_dim assert ( command.shape == (expected_dim,) and np.isfinite(command).all() ), f"Expected a finite goal command with dimension {expected_dim}, got shape {command.shape}" @@ -292,3 +440,11 @@ def reset(self, env_ids: torch.Tensor | None = None) -> None: else: for env_id in env_ids.reshape(-1).tolist(): self._goals.pop(env_id, None) + + +def _get_vlm_embodiment_adapter(name: str | None): + """Return the named action and prompt preset.""" + assert name == "droid", "Set vlm_embodiment_adapter='droid' or inject an action adapter and system_prompt" + from isaaclab_arena_vlm_agent.droid_adapter import DroidVLMAdapter + + return DroidVLMAdapter()