From 410789a78bd5f346b8094660881a47a71e05b93e Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Fri, 23 Oct 2026 21:10:31 -0700 Subject: [PATCH 01/11] Add VLM policies to experiment runner Signed-off-by: Xinjie Yao --- isaaclab_arena/embodiments/droid/droid.py | 14 ++ .../droid_pnp_vlm_experiment.yaml | 52 +++++ .../absolute_action_adapter.py | 20 +- isaaclab_arena_vlm_agent/droid_adapter.py | 67 ++++++ .../observation_adapter.py | 20 +- .../tests/test_droid_vlm_adapters.py | 2 +- .../test_droid_vlm_adapters_simulation.py | 10 +- .../tests/test_vlm_agent_policy.py | 12 +- .../tests/test_vlm_policy_configuration.py | 210 ++++++++++++++++++ ...est_vlm_policy_configuration_simulation.py | 69 ++++++ isaaclab_arena_vlm_agent/vlm_agent_policy.py | 191 +++++++++++++--- 11 files changed, 619 insertions(+), 48 deletions(-) create mode 100644 isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml create mode 100644 isaaclab_arena_vlm_agent/droid_adapter.py create mode 100644 isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py create mode 100644 isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index cb8637b033..4b816bbe0d 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -291,6 +291,20 @@ def _configure_newton_diff_ik(self) -> None: ) +@register_asset +class DroidAbsoluteDifferentialIKEmbodiment(DroidDifferentialIKEmbodiment): + """DROID differential IK accepting unscaled root-relative XYZ/XYZW pose targets.""" + + name = "droid_abs_ik" + + def get_action_cfg(self): + """Return absolute pose actions after applying the selected physics backend's IK settings.""" + cfg = super().get_action_cfg().copy() + cfg.arm_action.controller.use_relative_mode = False + cfg.arm_action.scale = 1.0 + return cfg + + @register_asset class DroidRelativeJointPositionEmbodiment(DroidEmbodimentBase): """Embodiment for the DROID setup with relative joint position action controller.""" diff --git a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml new file mode 100644 index 0000000000..700d763f59 --- /dev/null +++ b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml @@ -0,0 +1,52 @@ +# 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 --headless \ +# --experiment_config isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_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: &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: &builder + num_envs: 1 + language_instruction: Pick up the Rubik's cube and place it in the bowl. + policy: + type: isaaclab_arena_vlm_agent.vlm_agent_policy.VLMStepPolicy + endpoint: internal + vlm_embodiment_adapter: droid + observation_modality: image_state + rollout_limit: &rollout + num_episodes: 1 + + droid_pnp_vlm_chunk: + environment: *environment + environment_builder: *builder + policy: + type: isaaclab_arena_vlm_agent.vlm_agent_policy.VLMChunkPolicy + endpoint: internal + vlm_embodiment_adapter: droid + observation_modality: image_state + chunk_size: 3 + rollout_limit: *rollout + + droid_pnp_vlm_goal: + environment: *environment + environment_builder: *builder + policy: + type: isaaclab_arena_vlm_agent.vlm_agent_policy.VLMGoalPolicy + endpoint: internal + vlm_embodiment_adapter: droid + observation_modality: image_state + max_steps_per_goal: 30 + rollout_limit: *rollout diff --git a/isaaclab_arena_vlm_agent/absolute_action_adapter.py b/isaaclab_arena_vlm_agent/absolute_action_adapter.py index 14301b4112..47a16e800c 100644 --- a/isaaclab_arena_vlm_agent/absolute_action_adapter.py +++ b/isaaclab_arena_vlm_agent/absolute_action_adapter.py @@ -7,13 +7,31 @@ 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 AbsoluteStepActionAdapter(AgentActionAdapter): + """Decode a single absolute command in native controller order.""" + + response_model = Command + + def __init__(self): + self._command_indices = None + + def is_robot_action_space_supported(self, env) -> bool: + """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 AbsoluteChunkActionAdapter(ChunkActionAdapter): """Decode absolute command chunks; use chunk_size=1 for one prediction per step.""" diff --git a/isaaclab_arena_vlm_agent/droid_adapter.py b/isaaclab_arena_vlm_agent/droid_adapter.py new file mode 100644 index 0000000000..3d29939522 --- /dev/null +++ b/isaaclab_arena_vlm_agent/droid_adapter.py @@ -0,0 +1,67 @@ +# 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.pose_target_motion_generator import PoseTargetMotionGenerator + + +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 AbsoluteStepActionAdapter() + + def make_chunk_action_adapter(self, chunk_size: int) -> AbsoluteChunkActionAdapter: + """Create a decoder for the requested number of absolute pose commands.""" + return AbsoluteChunkActionAdapter(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 AbsoluteGoalActionAdapter(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 " + "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. Approach with the gripper " + "open, close around the object, lift it, move above the destination, and open to release. " + "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/observation_adapter.py b/isaaclab_arena_vlm_agent/observation_adapter.py index 1ba313caa6..449ba4ad76 100644 --- a/isaaclab_arena_vlm_agent/observation_adapter.py +++ b/isaaclab_arena_vlm_agent/observation_adapter.py @@ -92,21 +92,29 @@ 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 configured by an embodiment or the live environment.""" - def __init__(self, embodiment: EmbodimentBase): - observation_cfg = embodiment.get_observation_cfg() + def __init__(self, embodiment: EmbodimentBase | None = None): + self.camera_keys = None + if embodiment is not None: + self._set_camera_keys(embodiment.get_observation_cfg()) + + def _set_camera_keys(self, observation_cfg): + """Select RGB camera terms from the supplied observation configuration.""" camera_cfg = getattr(observation_cfg, 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" + assert camera_keys, "RGB modality requires at least one configured RGB observation" + self.camera_keys = camera_keys def extract_vlm_input_from_observation(self, env, observation, env_id: int): """Return configured RGB views for the selected environment.""" + if self.camera_keys is None: + self._set_camera_keys(env.unwrapped.cfg.observations) images = extract_rgb_images_from_camera_observations( observation[CAMERA_OBSERVATION_GROUP], self.camera_keys, env_id ) @@ -116,7 +124,7 @@ def extract_vlm_input_from_observation(self, env, observation, env_id: int): class ImageStateObservationAdapter(VLMObservationAdapter): """Include robot state followed by current RGB camera views.""" - def __init__(self, embodiment: EmbodimentBase): + def __init__(self, embodiment: EmbodimentBase | None = None): self.state_adapter = StateObservationAdapter() self.image_adapter = ImageObservationAdapter(embodiment) 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..ca45b963de 100644 --- a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py +++ b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py @@ -38,13 +38,13 @@ def inputs(monkeypatch): env = SimpleNamespace( unwrapped=SimpleNamespace( num_envs=2, - action_dim=8, device="cpu", 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), 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..04c8445207 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: 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..8ddca5b98d 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): @@ -234,7 +236,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 +267,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 +295,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..8428a7d8ca --- /dev/null +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py @@ -0,0 +1,210 @@ +# 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 {} + cfg = cfg_type(system_prompt="Return absolute targets.", model="test-model", **options) + cfg.observation_modality = modality + if cfg_type is VLMChunkPolicyCfg: + cfg.chunk_size = 2 + 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"})) + ) + env = SimpleNamespace( + unwrapped=SimpleNamespace( + num_envs=2, + device="cpu", + 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() + policy.close() + backend.client.close.assert_called_once() + assert not getattr(policy, "_chunks", {}) + assert not getattr(policy, "_goals", {}) + + +@pytest.mark.parametrize("policy_type,cfg_type", POLICIES) +def test_injected_backend_remains_caller_owned(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.client.close.assert_not_called() + + +@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() + + +def test_example_yaml_policy_selectors_and_round_trip(monkeypatch, tmp_path): + 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}, + ) + for policy_type, cfg_type in POLICIES: + 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] + experiment = loader.load_arena_experiment_from_config_file( + root / "isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml", + device="cpu", + overrides=[ + "runs.droid_pnp_vlm_goal.policy.max_steps_per_goal=5", + "runs.droid_pnp_vlm_chunk.policy.chunk_size=2", + ], + ) + assert len(experiment.runs) == 3 + for run, (policy_type, cfg_type) in zip(experiment.runs.values(), POLICIES, strict=True): + assert isinstance(run.policy, cfg_type) + assert run.environment_builder.device == "cpu" + assert registries.PolicyRegistry().get_policy_type_for_cfg(run.policy) is policy_type + assert experiment.runs["droid_pnp_vlm_goal"].policy.max_steps_per_goal == 5 + assert experiment.runs["droid_pnp_vlm_chunk"].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): + VLMChunkPolicyCfg(**options) 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..be8ef8068a --- /dev/null +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py @@ -0,0 +1,69 @@ +# 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 all three example strategies through typed loading and run execution.""" + +import pytest + +from isaaclab_arena.tests.utils.persistent_simulation_app import run_function_with_persistent_simulation_app + + +def _test_vlm_experiment(simulation_app, mode): + import json + import tempfile + from pathlib import Path + from unittest.mock import Mock, patch + + from isaaclab_arena.assets.registries import PolicyRegistry + 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_experiment.yaml" + run_name = f"droid_pnp_vlm_{mode}" + 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", + ], + ) + run = experiment.runs[run_name] + assert PolicyRegistry().get_policy_type_for_cfg(run.policy) is PolicyRegistry().get_policy(f"vlm_{mode}") + backend = Mock() + + def hold_current_pose(request): + content = request.messages[1]["content"] + assert any(item["type"] == "image_url" for item in content) + state = json.loads(content[1]["text"]) + pose = state["command_body_poses"]["arm_action"] + command = {"values": pose["position_root_m"] + pose["quaternion_root_xyzw"] + [0.0]} + if mode == "chunk": + assert f"exactly {run.policy.chunk_size} commands" in request.messages[0]["content"] + return {"commands": [command] * run.policy.chunk_size} + return command + + backend.infer.side_effect = hold_current_pose + 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 + if mode == "step": + assert backend.infer.call_count == 3 + elif mode == "chunk": + assert backend.infer.call_count == 1 + else: + assert 1 <= backend.infer.call_count <= 3 + backend.client.close.assert_called_once() + return True + + +@pytest.mark.with_cameras +@pytest.mark.parametrize("mode", ["step", "chunk", "goal"]) +def test_vlm_experiment(mode): + assert run_function_with_persistent_simulation_app(_test_vlm_experiment, enable_cameras=True, mode=mode) diff --git a/isaaclab_arena_vlm_agent/vlm_agent_policy.py b/isaaclab_arena_vlm_agent/vlm_agent_policy.py index 9671ed7b58..94a38e735e 100644 --- a/isaaclab_arena_vlm_agent/vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/vlm_agent_policy.py @@ -18,11 +18,17 @@ 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.observation_adapter import ( + ImageObservationAdapter, + ImageStateObservationAdapter, + StateObservationAdapter, + VLMObservationAdapter, +) from isaaclab_arena_vlm_agent.response_errors import InvalidAgentResponseError if TYPE_CHECKING: @@ -33,8 +39,8 @@ 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 +52,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: str = "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'" + assert self.observation_modality in {"image", "state", "image_state"} 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 = 1 + """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 +154,60 @@ 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: Caller-owned backend, or None to create a client owned by this policy. + 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 = { + "image": ImageObservationAdapter, + "state": StateObservationAdapter, + "image_state": ImageStateObservationAdapter, + }[config.observation_modality] + 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 + self._owns_inference_backend = inference_backend is None + self._closed = False + # 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: + """Discard pending commands and close only a policy-owned inference client.""" + if not self._closed: + self.reset() + if self._owns_inference_backend: + self._inference_backend.client.close() + self._closed = True def get_action(self, env: gym.Env, observation: dict[str, Any]) -> torch.Tensor: """Return a batch of native actions using the policy execution strategy. @@ -147,7 +256,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 +294,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 +322,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 +365,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 +407,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 +421,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() From e132f434acd615fd95f8baea35822d69ae4ad47c Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Fri, 23 Oct 2026 21:13:23 -0700 Subject: [PATCH 02/11] Use registry_name for the DROID IK embodiment Signed-off-by: Xinjie Yao --- isaaclab_arena/embodiments/droid/droid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 4b816bbe0d..5e31290f20 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -295,7 +295,7 @@ def _configure_newton_diff_ik(self) -> None: class DroidAbsoluteDifferentialIKEmbodiment(DroidDifferentialIKEmbodiment): """DROID differential IK accepting unscaled root-relative XYZ/XYZW pose targets.""" - name = "droid_abs_ik" + registry_name = "droid_abs_ik" def get_action_cfg(self): """Return absolute pose actions after applying the selected physics backend's IK settings.""" From bc3836e7b6c022952fcd99ff59a9ea8353f9d64e Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Sun, 25 Oct 2026 23:01:53 -0700 Subject: [PATCH 03/11] Add live camera calibration to VLM policy inputs Signed-off-by: Xinjie Yao --- isaaclab_arena/embodiments/droid/droid.py | 3 + .../embodiments/droid/observations.py | 2 +- .../droid_pnp_vlm_experiment.yaml | 3 + isaaclab_arena_vlm_agent/droid_adapter.py | 6 ++ isaaclab_arena_vlm_agent/image_utils.py | 6 +- .../observation_adapter.py | 56 ++++++++++++- .../tests/test_droid_vlm_adapters.py | 19 ++++- .../tests/test_image_utils.py | 4 +- .../tests/test_observation_adapter.py | 82 +++++++++++++++++-- .../tests/test_vlm_agent_policy.py | 21 +++++ .../tests/test_vlm_policy_configuration.py | 18 +++- ...est_vlm_policy_configuration_simulation.py | 49 ++++++++++- isaaclab_arena_vlm_agent/vlm_agent_policy.py | 15 +++- 13 files changed, 263 insertions(+), 21 deletions(-) diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 5e31290f20..4e9ea083b4 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -583,6 +583,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, @@ -596,6 +597,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, @@ -609,6 +611,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_environments/experiment_configs/droid_pnp_vlm_experiment.yaml b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml index 700d763f59..3940b27bc7 100644 --- a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml +++ b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml @@ -24,6 +24,7 @@ runs: 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: &rollout @@ -35,6 +36,7 @@ runs: 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 @@ -46,6 +48,7 @@ runs: 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 diff --git a/isaaclab_arena_vlm_agent/droid_adapter.py b/isaaclab_arena_vlm_agent/droid_adapter.py index 3d29939522..2867861709 100644 --- a/isaaclab_arena_vlm_agent/droid_adapter.py +++ b/isaaclab_arena_vlm_agent/droid_adapter.py @@ -53,6 +53,12 @@ def get_system_prompt(self, action_adapter: AgentActionAdapter) -> str: + 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. Approach with the gripper " "open, close around the object, lift it, move above the destination, and open to release. " diff --git a/isaaclab_arena_vlm_agent/image_utils.py b/isaaclab_arena_vlm_agent/image_utils.py index 1cf7eb6921..51d015be18 100644 --- a/isaaclab_arena_vlm_agent/image_utils.py +++ b/isaaclab_arena_vlm_agent/image_utils.py @@ -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,4 @@ 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 diff --git a/isaaclab_arena_vlm_agent/observation_adapter.py b/isaaclab_arena_vlm_agent/observation_adapter.py index 449ba4ad76..c85b7402f7 100644 --- a/isaaclab_arena_vlm_agent/observation_adapter.py +++ b/isaaclab_arena_vlm_agent/observation_adapter.py @@ -9,6 +9,7 @@ import numpy as np from abc import ABC, abstractmethod +from scipy.spatial.transform import Rotation from typing import TYPE_CHECKING, Any from isaaclab_arena_vlm_agent.image_utils import extract_rgb_images_from_camera_observations @@ -112,13 +113,60 @@ def _set_camera_keys(self, observation_cfg): self.camera_keys = camera_keys def extract_vlm_input_from_observation(self, env, observation, env_id: int): - """Return configured RGB views for the selected environment.""" + """Return RGB views and optical-frame camera calibration for one environment.""" if self.camera_keys is None: self._set_camera_keys(env.unwrapped.cfg.observations) images = extract_rgb_images_from_camera_observations( observation[CAMERA_OBSERVATION_GROUP], self.camera_keys, env_id ) - return images, {} + 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 images, {"camera_geometry": geometry} class ImageStateObservationAdapter(VLMObservationAdapter): @@ -131,5 +179,5 @@ def __init__(self, embodiment: EmbodimentBase | None = None): 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} 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 ca45b963de..3598ae9dbf 100644 --- a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py +++ b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py @@ -30,15 +30,32 @@ 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, 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} 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..223ed20cd2 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} @@ -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(): @@ -154,3 +174,55 @@ 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) 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 8ddca5b98d..d708b910a3 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py @@ -133,6 +133,27 @@ 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]]) + policy.observation_adapter.extract_vlm_input_from_observation = lambda *args: ( + {"front_view": pixels}, + {"camera_geometry": {"front_view": {"image_size_px": [17, 9], "intrinsic_matrix": intrinsics.tolist()}}}, + ) + content = policy._request_content(env, observation, 0) + 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 diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py index 8428a7d8ca..f8b104a00a 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py @@ -71,12 +71,27 @@ def test_configured_execution_and_partial_reset(monkeypatch, policy_type, cfg_ty 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"})) + 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) ), @@ -175,6 +190,7 @@ def test_example_yaml_policy_selectors_and_round_trip(monkeypatch, tmp_path): assert len(experiment.runs) == 3 for run, (policy_type, cfg_type) in zip(experiment.runs.values(), POLICIES, strict=True): 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 assert experiment.runs["droid_pnp_vlm_goal"].policy.max_steps_per_goal == 5 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 index be8ef8068a..842cd805e7 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py @@ -11,12 +11,19 @@ def _test_vlm_experiment(simulation_app, mode): + 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 @@ -36,25 +43,63 @@ def _test_vlm_experiment(simulation_app, mode): run = experiment.runs[run_name] assert PolicyRegistry().get_policy_type_for_cfg(run.policy) is PolicyRegistry().get_policy(f"vlm_{mode}") backend = Mock() + wrist_positions = [] + wrist_offset = DroidCameraCfg().wrist_camera.offset - def hold_current_pose(request): + 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 if mode == "chunk": assert f"exactly {run.policy.chunk_size} commands" in request.messages[0]["content"] return {"commands": [command] * run.policy.chunk_size} return command - backend.infer.side_effect = hold_current_pose + 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 if mode == "step": assert backend.infer.call_count == 3 + assert np.linalg.norm(np.array(wrist_positions[-1]) - wrist_positions[0]) > 1e-4 elif mode == "chunk": assert backend.infer.call_count == 1 else: diff --git a/isaaclab_arena_vlm_agent/vlm_agent_policy.py b/isaaclab_arena_vlm_agent/vlm_agent_policy.py index 94a38e735e..d91ceb6b6a 100644 --- a/isaaclab_arena_vlm_agent/vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/vlm_agent_policy.py @@ -240,15 +240,26 @@ 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) + 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", {}): + geometry = prompt_values["camera_geometry"][camera_name] + intrinsics = np.array(geometry["intrinsic_matrix"], dtype=float) + intrinsics[0] *= size[0] / image.shape[1] + intrinsics[1] *= size[1] / image.shape[0] + geometry["intrinsic_matrix"] = intrinsics.tolist() + geometry["image_size_px"] = list(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 From a0d0463436d88bb4b18277c2964dbadd91c5b37b Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 13:45:51 -0700 Subject: [PATCH 04/11] Address VLM policy review feedback Signed-off-by: Xinjie Yao --- isaaclab_arena/embodiments/droid/droid.py | 19 +++-- .../arena_experiment_config_loader.py | 3 + isaaclab_arena/inference/backend.py | 4 + .../absolute_action_adapter.py | 36 +++------ isaaclab_arena_vlm_agent/droid_adapter.py | 31 ++++++- .../observation_adapter.py | 35 +++----- isaaclab_arena_vlm_agent/tests/conftest.py | 8 ++ .../tests/test_droid_vlm_adapters.py | 2 +- .../test_droid_vlm_adapters_simulation.py | 28 +++++++ .../tests/test_observation_adapter.py | 22 ++++- .../tests/test_vlm_agent_policy.py | 9 ++- .../tests/test_vlm_policy_configuration.py | 81 ++++++++++++++++--- ...est_vlm_policy_configuration_simulation.py | 43 +++++----- isaaclab_arena_vlm_agent/vlm_agent_policy.py | 43 ++++++---- 14 files changed, 254 insertions(+), 110 deletions(-) create mode 100644 isaaclab_arena_vlm_agent/tests/conftest.py diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 4e9ea083b4..ce3abe23df 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -297,12 +297,19 @@ class DroidAbsoluteDifferentialIKEmbodiment(DroidDifferentialIKEmbodiment): registry_name = "droid_abs_ik" - def get_action_cfg(self): - """Return absolute pose actions after applying the selected physics backend's IK settings.""" - cfg = super().get_action_cfg().copy() - cfg.arm_action.controller.use_relative_mode = False - cfg.arm_action.scale = 1.0 - return cfg + 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 diff --git a/isaaclab_arena/evaluation/arena_experiment_config_loader.py b/isaaclab_arena/evaluation/arena_experiment_config_loader.py index 92bc442e68..0826e58c3e 100644 --- a/isaaclab_arena/evaluation/arena_experiment_config_loader.py +++ b/isaaclab_arena/evaluation/arena_experiment_config_loader.py @@ -107,6 +107,9 @@ def _resolve_policy_cfg_type_from_name_or_class_path(policy_name_or_class_path: from isaaclab_arena.assets.registries import PolicyRegistry registry = PolicyRegistry() + # Optional policy packages are imported only when their short names are selected. + if policy_name_or_class_path in {"vlm_step", "vlm_chunk", "vlm_goal"}: + import_module("isaaclab_arena_vlm_agent.vlm_agent_policy") if registry.is_registered(policy_name_or_class_path): policy_type = registry.get_policy(policy_name_or_class_path) else: 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_vlm_agent/absolute_action_adapter.py b/isaaclab_arena_vlm_agent/absolute_action_adapter.py index 47a16e800c..8399addaae 100644 --- a/isaaclab_arena_vlm_agent/absolute_action_adapter.py +++ b/isaaclab_arena_vlm_agent/absolute_action_adapter.py @@ -14,7 +14,7 @@ from isaaclab_arena_vlm_agent.response_errors import InvalidAgentResponseError -class AbsoluteStepActionAdapter(AgentActionAdapter): +class _AbsoluteCommandAdapter(AgentActionAdapter): """Decode a single absolute command in native controller order.""" response_model = Command @@ -32,19 +32,18 @@ def decode_command(self, response: Command, arena_world, env_id: int) -> np.ndar return _decode_absolute_command(response, self._command_indices) -class AbsoluteChunkActionAdapter(ChunkActionAdapter): - """Decode absolute command chunks; use chunk_size=1 for one prediction per step.""" +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): - super().__init__(chunk_size) - 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 + _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.""" @@ -52,27 +51,18 @@ def decode_command(self, response: CommandChunk, arena_world, env_id: int) -> np 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 index 2867861709..1b2cb17405 100644 --- a/isaaclab_arena_vlm_agent/droid_adapter.py +++ b/isaaclab_arena_vlm_agent/droid_adapter.py @@ -11,19 +11,44 @@ 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 AbsoluteStepActionAdapter() + return _DroidStepActionAdapter() def make_chunk_action_adapter(self, chunk_size: int) -> AbsoluteChunkActionAdapter: """Create a decoder for the requested number of absolute pose commands.""" - return AbsoluteChunkActionAdapter(chunk_size) + return _DroidChunkActionAdapter(chunk_size) def make_goal_action_adapter( self, position_tolerance_m: float, rotation_tolerance_rad: float @@ -34,7 +59,7 @@ def make_goal_action_adapter( position_tolerance_m=position_tolerance_m, rotation_tolerance_rad=rotation_tolerance_rad, ) - return AbsoluteGoalActionAdapter(generator) + return _DroidGoalActionAdapter(generator) def get_system_prompt(self, action_adapter: AgentActionAdapter) -> str: """Describe the DROID pose contract and the selected response schema.""" diff --git a/isaaclab_arena_vlm_agent/observation_adapter.py b/isaaclab_arena_vlm_agent/observation_adapter.py index c85b7402f7..0589520f35 100644 --- a/isaaclab_arena_vlm_agent/observation_adapter.py +++ b/isaaclab_arena_vlm_agent/observation_adapter.py @@ -10,13 +10,10 @@ import numpy as np from abc import ABC, abstractmethod from scipy.spatial.transform import Rotation -from typing import TYPE_CHECKING, Any +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.""" @@ -93,16 +90,12 @@ def extract_vlm_input_from_observation(self, env, observation, env_id: int): class ImageObservationAdapter(VLMObservationAdapter): - """Include RGB views configured by an embodiment or the live environment.""" - - def __init__(self, embodiment: EmbodimentBase | None = None): - self.camera_keys = None - if embodiment is not None: - self._set_camera_keys(embodiment.get_observation_cfg()) + """Include RGB views and calibration from the live environment.""" - def _set_camera_keys(self, observation_cfg): - """Select RGB camera terms from the supplied observation configuration.""" - 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" camera_keys = tuple( name @@ -110,17 +103,7 @@ def _set_camera_keys(self, observation_cfg): if hasattr(term, "func") and term.params.get("data_type") == "rgb" ) assert camera_keys, "RGB modality requires at least one configured RGB observation" - self.camera_keys = camera_keys - - def extract_vlm_input_from_observation(self, env, observation, env_id: int): - """Return RGB views and optical-frame camera calibration for one environment.""" - if self.camera_keys is None: - self._set_camera_keys(env.unwrapped.cfg.observations) - images = extract_rgb_images_from_camera_observations( - observation[CAMERA_OBSERVATION_GROUP], self.camera_keys, env_id - ) - base = env.unwrapped - camera_cfg = getattr(base.cfg.observations, CAMERA_OBSERVATION_GROUP) + images = extract_rgb_images_from_camera_observations(observation[CAMERA_OBSERVATION_GROUP], camera_keys, env_id) roots = {} for name in base.action_manager.active_terms: asset_name = base.action_manager.get_term(name).cfg.asset_name @@ -172,9 +155,9 @@ def extract_vlm_input_from_observation(self, env, observation, env_id: int): class ImageStateObservationAdapter(VLMObservationAdapter): """Include robot state followed by current RGB camera views.""" - def __init__(self, embodiment: EmbodimentBase | None = None): + 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.""" 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 3598ae9dbf..47ff7c2a73 100644 --- a/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py +++ b/isaaclab_arena_vlm_agent/tests/test_droid_vlm_adapters.py @@ -97,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 04c8445207..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 @@ -119,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 = [] @@ -205,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_observation_adapter.py b/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py index 223ed20cd2..ec12779d8b 100644 --- a/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py +++ b/isaaclab_arena_vlm_agent/tests/test_observation_adapter.py @@ -109,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]} @@ -149,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(): @@ -164,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( @@ -226,3 +226,17 @@ def test_camera_geometry_rejects_stale_pose_configuration(): 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 d708b910a3..6b4b8cb02b 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_agent_policy.py @@ -140,11 +140,12 @@ def test_request_intrinsics_match_encoded_pixels(backend, max_size, expected_siz 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]]) - policy.observation_adapter.extract_vlm_input_from_observation = lambda *args: ( - {"front_view": pixels}, - {"camera_geometry": {"front_view": {"image_size_px": [17, 9], "intrinsic_matrix": intrinsics.tolist()}}}, - ) + 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"]) diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py index f8b104a00a..0b951e21fe 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py @@ -42,10 +42,10 @@ def gripper_measurements(monkeypatch): 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 - if cfg_type is VLMChunkPolicyCfg: - cfg.chunk_size = 2 return cfg @@ -125,10 +125,7 @@ def test_configured_execution_and_partial_reset(monkeypatch, policy_type, cfg_ty assert backend.infer.call_count == 4 assert constructor.call_args.kwargs["model"] == "test-model" policy.close() - policy.close() - backend.client.close.assert_called_once() - assert not getattr(policy, "_chunks", {}) - assert not getattr(policy, "_goals", {}) + backend.close.assert_called_once() @pytest.mark.parametrize("policy_type,cfg_type", POLICIES) @@ -139,7 +136,7 @@ def test_injected_backend_remains_caller_owned(monkeypatch, policy_type, cfg_typ policy = policy_type(_config(cfg_type), inference_backend=backend) policy.close() constructor.assert_not_called() - backend.client.close.assert_not_called() + backend.close.assert_not_called() @pytest.mark.parametrize("policy_type,cfg_type", POLICIES) @@ -222,5 +219,71 @@ def test_preset_supplies_prompt_and_strategy(monkeypatch, policy_type, cfg_type) "options", [{"vlm_embodiment_adapter": "unknown"}, {"observation_modality": "unknown"}, {"chunk_size": 0}] ) def test_invalid_flat_settings(options): - with pytest.raises(AssertionError): - VLMChunkPolicyCfg(**options) + with pytest.raises((AssertionError, ValueError)): + VLMChunkPolicyCfg(**{"chunk_size": 2, **options}) + + +@pytest.mark.parametrize("name", ["vlm_step", "vlm_chunk", "vlm_goal"]) +def test_short_policy_name_is_discovered_in_fresh_process(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({name!r}); " + "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 index 842cd805e7..1408025785 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration_simulation.py @@ -3,14 +3,16 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Verify all three example strategies through typed loading and run execution.""" +"""Verify goal-policy execution and moving-camera calibration in an isolated simulator.""" + +import sys import pytest -from isaaclab_arena.tests.utils.persistent_simulation_app import run_function_with_persistent_simulation_app +from isaaclab_arena.tests.utils.subprocess import run_subprocess -def _test_vlm_experiment(simulation_app, mode): +def _test_vlm_experiment(simulation_app): import base64 import io import json @@ -30,7 +32,7 @@ def _test_vlm_experiment(simulation_app, mode): from isaaclab_arena.tests.utils.constants import TestConstants config_path = Path(TestConstants.arena_environments_dir) / "experiment_configs/droid_pnp_vlm_experiment.yaml" - run_name = f"droid_pnp_vlm_{mode}" + run_name = "droid_pnp_vlm_goal" experiment = load_arena_experiment_from_config_file( config_path, device="cuda:0", @@ -38,10 +40,11 @@ def _test_vlm_experiment(simulation_app, mode): 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(f"vlm_{mode}") + 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 @@ -87,9 +90,6 @@ def move_nearby(request): wrist_positions.append(wrist["position_world_m"]) command = {"values": pose["position_root_m"] + pose["quaternion_root_xyzw"] + [0.0]} command["values"][0] += 0.01 - if mode == "chunk": - assert f"exactly {run.policy.chunk_size} commands" in request.messages[0]["content"] - return {"commands": [command] * run.policy.chunk_size} return command backend.infer.side_effect = move_nearby @@ -97,18 +97,23 @@ def move_nearby(request): 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 - if mode == "step": - assert backend.infer.call_count == 3 - assert np.linalg.norm(np.array(wrist_positions[-1]) - wrist_positions[0]) > 1e-4 - elif mode == "chunk": - assert backend.infer.call_count == 1 - else: - assert 1 <= backend.infer.call_count <= 3 - backend.client.close.assert_called_once() + 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.parametrize("mode", ["step", "chunk", "goal"]) -def test_vlm_experiment(mode): - assert run_function_with_persistent_simulation_app(_test_vlm_experiment, enable_cameras=True, mode=mode) +@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 d91ceb6b6a..c00fc40d71 100644 --- a/isaaclab_arena_vlm_agent/vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/vlm_agent_policy.py @@ -13,7 +13,9 @@ 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 @@ -35,6 +37,23 @@ 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.""" @@ -55,7 +74,7 @@ class VLMPolicyCfg(PolicyCfg): vlm_embodiment_adapter: str | None = None """Named action and prompt preset; droid supports the droid_abs_ik embodiment.""" - observation_modality: str = "image_state" + observation_modality: ObservationModality = ObservationModality.image_state """One of image, state, or image_state; ignored when an observation adapter is injected.""" endpoint: str | None = None @@ -82,7 +101,7 @@ def __post_init__(self): 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'" - assert self.observation_modality in {"image", "state", "image_state"} + 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 @@ -110,7 +129,7 @@ class VLMChunkPolicyCfg(VLMPolicyCfg): vlm_embodiment_adapter: str | None = "droid" """Named action and prompt preset.""" - chunk_size: int = 1 + chunk_size: int = field(kw_only=True) """Number of commands per prediction; also used by the default prompt.""" def __post_init__(self): @@ -168,11 +187,7 @@ def __init__( """ super().__init__(config) if observation_adapter is None: - adapter_type = { - "image": ImageObservationAdapter, - "state": StateObservationAdapter, - "image_state": ImageStateObservationAdapter, - }[config.observation_modality] + adapter_type = ObservationModality(config.observation_modality).get_adapter_type() observation_adapter = adapter_type() assert isinstance(observation_adapter, VLMObservationAdapter) assert isinstance(action_adapter, AgentActionAdapter) @@ -186,7 +201,6 @@ def __init__( ) self.task_description = None self._owns_inference_backend = inference_backend is None - self._closed = False # Construct the client last so invalid adapters or schemas cannot leak a connection. self._inference_backend = ( inference_backend @@ -202,12 +216,9 @@ def __init__( ) def close(self) -> None: - """Discard pending commands and close only a policy-owned inference client.""" - if not self._closed: - self.reset() - if self._owns_inference_backend: - self._inference_backend.client.close() - self._closed = True + """Close the inference backend when this policy created it.""" + if self._owns_inference_backend: + 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. @@ -240,6 +251,8 @@ 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) From d468e06a85d568dd0cde5b1dbe08b794100c8034 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 14:08:37 -0700 Subject: [PATCH 05/11] Discover optional policies through entry points Signed-off-by: Xinjie Yao --- .../arena_experiment_config_loader.py | 19 +++++++----- .../test_arena_experiment_config_loader.py | 31 +++++++++++++++++++ pyproject.toml | 6 ++++ 3 files changed, 48 insertions(+), 8 deletions(-) diff --git a/isaaclab_arena/evaluation/arena_experiment_config_loader.py b/isaaclab_arena/evaluation/arena_experiment_config_loader.py index 0826e58c3e..9392e519b2 100644 --- a/isaaclab_arena/evaluation/arena_experiment_config_loader.py +++ b/isaaclab_arena/evaluation/arena_experiment_config_loader.py @@ -10,6 +10,7 @@ import json from dataclasses import replace from importlib import import_module +from importlib.metadata import entry_points from pathlib import Path from typing import TYPE_CHECKING, Any @@ -103,19 +104,21 @@ def _registered_environment_cfg_types() -> dict[str, type[ArenaEnvironmentCfg]]: def _resolve_policy_cfg_type_from_name_or_class_path(policy_name_or_class_path: str) -> type[PolicyCfg]: - """Return the config type for a registered policy name or dotted class path.""" + """Return the config type for a registered name, policy entry point, or dotted class path.""" from isaaclab_arena.assets.registries import PolicyRegistry registry = PolicyRegistry() - # Optional policy packages are imported only when their short names are selected. - if policy_name_or_class_path in {"vlm_step", "vlm_chunk", "vlm_goal"}: - import_module("isaaclab_arena_vlm_agent.vlm_agent_policy") if registry.is_registered(policy_name_or_class_path): policy_type = registry.get_policy(policy_name_or_class_path) - else: - assert ( - "." in policy_name_or_class_path - ), f"Policy type must be a registered name or dotted Python class path, got {policy_name_or_class_path!r}" + elif "." in policy_name_or_class_path: module_path, class_name = policy_name_or_class_path.rsplit(".", 1) policy_type = getattr(import_module(module_path), class_name) + else: + matches = entry_points(group="isaaclab_arena.policies", name=policy_name_or_class_path) + assert len(matches) == 1, ( + f"Expected one policy entry point for {policy_name_or_class_path!r}, found {len(matches)}. " + "Install the policy package or use a dotted Python class path." + ) + (entry_point,) = matches + policy_type = entry_point.load() return registry.get_policy_cfg_type(policy_type) diff --git a/isaaclab_arena/tests/test_arena_experiment_config_loader.py b/isaaclab_arena/tests/test_arena_experiment_config_loader.py index ba1f01c613..a5cada17eb 100644 --- a/isaaclab_arena/tests/test_arena_experiment_config_loader.py +++ b/isaaclab_arena/tests/test_arena_experiment_config_loader.py @@ -387,3 +387,34 @@ def test_unknown_experiment_config_file_type_is_rejected(tmp_path): with pytest.raises(AssertionError, match="must use .json, .yaml, or .yml"): validate_experiment_config_path(config_path) + + +def test_policy_config_type_resolves_from_package_entry_point(monkeypatch): + from importlib.metadata import EntryPoint, EntryPoints + from unittest.mock import Mock + + discovered = Mock( + return_value=EntryPoints([ + EntryPoint( + name="external_policy", + value="isaaclab_arena.policy.zero_action_policy:ZeroActionPolicy", + group="isaaclab_arena.policies", + ) + ]) + ) + monkeypatch.setattr(arena_experiment_config_loader, "entry_points", discovered) + cfg_type = arena_experiment_config_loader._resolve_policy_cfg_type_from_name_or_class_path("external_policy") + assert cfg_type is ZeroActionPolicyCfg + discovered.assert_called_once_with(group="isaaclab_arena.policies", name="external_policy") + + +@pytest.mark.parametrize("count", [0, 2]) +def test_policy_entry_point_must_be_unambiguous(monkeypatch, count): + from unittest.mock import Mock + + entries = [Mock() for _ in range(count)] + monkeypatch.setattr(arena_experiment_config_loader, "entry_points", Mock(return_value=entries)) + with pytest.raises(AssertionError, match=f"found {count}"): + arena_experiment_config_loader._resolve_policy_cfg_type_from_name_or_class_path("external_policy") + for entry in entries: + entry.load.assert_not_called() diff --git a/pyproject.toml b/pyproject.toml index 522a6af9fa..2241d9394e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,12 @@ dependencies = [ "pandas==2.2.3", ] +# Optional policies are imported only when their entry-point name is selected. +[project.entry-points."isaaclab_arena.policies"] +vlm_step = "isaaclab_arena_vlm_agent.vlm_agent_policy:VLMStepPolicy" +vlm_chunk = "isaaclab_arena_vlm_agent.vlm_agent_policy:VLMChunkPolicy" +vlm_goal = "isaaclab_arena_vlm_agent.vlm_agent_policy:VLMGoalPolicy" + [project.optional-dependencies] dev = [ "jupyter", From da4a9886e446a96fdc3d3800406dd842e182d9d9 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 14:32:29 -0700 Subject: [PATCH 06/11] Use direct class paths for optional policies Signed-off-by: Xinjie Yao --- .../arena_experiment_config_loader.py | 17 ++++------ .../test_arena_experiment_config_loader.py | 32 ++----------------- .../tests/test_vlm_policy_configuration.py | 6 ++-- pyproject.toml | 6 ---- 4 files changed, 12 insertions(+), 49 deletions(-) diff --git a/isaaclab_arena/evaluation/arena_experiment_config_loader.py b/isaaclab_arena/evaluation/arena_experiment_config_loader.py index 9392e519b2..1db2f5d48a 100644 --- a/isaaclab_arena/evaluation/arena_experiment_config_loader.py +++ b/isaaclab_arena/evaluation/arena_experiment_config_loader.py @@ -10,7 +10,6 @@ import json from dataclasses import replace from importlib import import_module -from importlib.metadata import entry_points from pathlib import Path from typing import TYPE_CHECKING, Any @@ -104,21 +103,17 @@ def _registered_environment_cfg_types() -> dict[str, type[ArenaEnvironmentCfg]]: def _resolve_policy_cfg_type_from_name_or_class_path(policy_name_or_class_path: str) -> type[PolicyCfg]: - """Return the config type for a registered name, policy entry point, or dotted class path.""" + """Return the config type for a registered policy name or dotted class path.""" from isaaclab_arena.assets.registries import PolicyRegistry registry = PolicyRegistry() if registry.is_registered(policy_name_or_class_path): policy_type = registry.get_policy(policy_name_or_class_path) - elif "." in policy_name_or_class_path: - module_path, class_name = policy_name_or_class_path.rsplit(".", 1) - policy_type = getattr(import_module(module_path), class_name) else: - matches = entry_points(group="isaaclab_arena.policies", name=policy_name_or_class_path) - assert len(matches) == 1, ( - f"Expected one policy entry point for {policy_name_or_class_path!r}, found {len(matches)}. " - "Install the policy package or use a dotted Python class path." + assert "." in policy_name_or_class_path, ( + f"Policy {policy_name_or_class_path!r} is not registered. " + "Use a dotted Python class path to import the policy." ) - (entry_point,) = matches - policy_type = entry_point.load() + module_path, class_name = policy_name_or_class_path.rsplit(".", 1) + policy_type = getattr(import_module(module_path), class_name) return registry.get_policy_cfg_type(policy_type) diff --git a/isaaclab_arena/tests/test_arena_experiment_config_loader.py b/isaaclab_arena/tests/test_arena_experiment_config_loader.py index a5cada17eb..417d27bbd9 100644 --- a/isaaclab_arena/tests/test_arena_experiment_config_loader.py +++ b/isaaclab_arena/tests/test_arena_experiment_config_loader.py @@ -389,32 +389,6 @@ def test_unknown_experiment_config_file_type_is_rejected(tmp_path): validate_experiment_config_path(config_path) -def test_policy_config_type_resolves_from_package_entry_point(monkeypatch): - from importlib.metadata import EntryPoint, EntryPoints - from unittest.mock import Mock - - discovered = Mock( - return_value=EntryPoints([ - EntryPoint( - name="external_policy", - value="isaaclab_arena.policy.zero_action_policy:ZeroActionPolicy", - group="isaaclab_arena.policies", - ) - ]) - ) - monkeypatch.setattr(arena_experiment_config_loader, "entry_points", discovered) - cfg_type = arena_experiment_config_loader._resolve_policy_cfg_type_from_name_or_class_path("external_policy") - assert cfg_type is ZeroActionPolicyCfg - discovered.assert_called_once_with(group="isaaclab_arena.policies", name="external_policy") - - -@pytest.mark.parametrize("count", [0, 2]) -def test_policy_entry_point_must_be_unambiguous(monkeypatch, count): - from unittest.mock import Mock - - entries = [Mock() for _ in range(count)] - monkeypatch.setattr(arena_experiment_config_loader, "entry_points", Mock(return_value=entries)) - with pytest.raises(AssertionError, match=f"found {count}"): - arena_experiment_config_loader._resolve_policy_cfg_type_from_name_or_class_path("external_policy") - for entry in entries: - entry.load.assert_not_called() +def test_unregistered_policy_name_requires_a_class_path(): + with pytest.raises(AssertionError, match="Use a dotted Python class path"): + arena_experiment_config_loader._resolve_policy_cfg_type_from_name_or_class_path("unregistered_policy") diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py index 0b951e21fe..66a19b2ae7 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py @@ -223,8 +223,8 @@ def test_invalid_flat_settings(options): VLMChunkPolicyCfg(**{"chunk_size": 2, **options}) -@pytest.mark.parametrize("name", ["vlm_step", "vlm_chunk", "vlm_goal"]) -def test_short_policy_name_is_discovered_in_fresh_process(name): +@pytest.mark.parametrize("class_name", ["VLMStepPolicy", "VLMChunkPolicy", "VLMGoalPolicy"]) +def test_policy_class_path_loads_in_fresh_process(class_name): import subprocess import sys @@ -237,7 +237,7 @@ def test_short_policy_name_is_discovered_in_fresh_process(name): "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({name!r}); " + f"cfg = resolve('isaaclab_arena_vlm_agent.vlm_agent_policy.{class_name}'); " "assert cfg.__module__ == 'isaaclab_arena_vlm_agent.vlm_agent_policy'" ), ], diff --git a/pyproject.toml b/pyproject.toml index 2241d9394e..522a6af9fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,12 +40,6 @@ dependencies = [ "pandas==2.2.3", ] -# Optional policies are imported only when their entry-point name is selected. -[project.entry-points."isaaclab_arena.policies"] -vlm_step = "isaaclab_arena_vlm_agent.vlm_agent_policy:VLMStepPolicy" -vlm_chunk = "isaaclab_arena_vlm_agent.vlm_agent_policy:VLMChunkPolicy" -vlm_goal = "isaaclab_arena_vlm_agent.vlm_agent_policy:VLMGoalPolicy" - [project.optional-dependencies] dev = [ "jupyter", From 37ed6dc5a855ea82c8f5494c32ef702ab282930e Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 14:35:01 -0700 Subject: [PATCH 07/11] Drop unnecessary experiment loader changes Signed-off-by: Xinjie Yao --- .../evaluation/arena_experiment_config_loader.py | 7 +++---- .../tests/test_arena_experiment_config_loader.py | 5 ----- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/isaaclab_arena/evaluation/arena_experiment_config_loader.py b/isaaclab_arena/evaluation/arena_experiment_config_loader.py index 1db2f5d48a..92bc442e68 100644 --- a/isaaclab_arena/evaluation/arena_experiment_config_loader.py +++ b/isaaclab_arena/evaluation/arena_experiment_config_loader.py @@ -110,10 +110,9 @@ def _resolve_policy_cfg_type_from_name_or_class_path(policy_name_or_class_path: if registry.is_registered(policy_name_or_class_path): policy_type = registry.get_policy(policy_name_or_class_path) else: - assert "." in policy_name_or_class_path, ( - f"Policy {policy_name_or_class_path!r} is not registered. " - "Use a dotted Python class path to import the policy." - ) + assert ( + "." in policy_name_or_class_path + ), f"Policy type must be a registered name or dotted Python class path, got {policy_name_or_class_path!r}" module_path, class_name = policy_name_or_class_path.rsplit(".", 1) policy_type = getattr(import_module(module_path), class_name) return registry.get_policy_cfg_type(policy_type) diff --git a/isaaclab_arena/tests/test_arena_experiment_config_loader.py b/isaaclab_arena/tests/test_arena_experiment_config_loader.py index 417d27bbd9..ba1f01c613 100644 --- a/isaaclab_arena/tests/test_arena_experiment_config_loader.py +++ b/isaaclab_arena/tests/test_arena_experiment_config_loader.py @@ -387,8 +387,3 @@ def test_unknown_experiment_config_file_type_is_rejected(tmp_path): with pytest.raises(AssertionError, match="must use .json, .yaml, or .yml"): validate_experiment_config_path(config_path) - - -def test_unregistered_policy_name_requires_a_class_path(): - with pytest.raises(AssertionError, match="Use a dotted Python class path"): - arena_experiment_config_loader._resolve_policy_cfg_type_from_name_or_class_path("unregistered_policy") From b84b3ab88bc07447f2d3c6d12ea047f574a2fc47 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 14:45:58 -0700 Subject: [PATCH 08/11] Simplify VLM ownership and generalize DROID prompt Signed-off-by: Xinjie Yao --- isaaclab_arena_vlm_agent/droid_adapter.py | 3 +-- .../tests/test_vlm_policy_configuration.py | 4 ++-- isaaclab_arena_vlm_agent/vlm_agent_policy.py | 8 +++----- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/isaaclab_arena_vlm_agent/droid_adapter.py b/isaaclab_arena_vlm_agent/droid_adapter.py index 1b2cb17405..0892b09365 100644 --- a/isaaclab_arena_vlm_agent/droid_adapter.py +++ b/isaaclab_arena_vlm_agent/droid_adapter.py @@ -85,8 +85,7 @@ def get_system_prompt(self, action_adapter: AgentActionAdapter) -> str: "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. Approach with the gripper " - "open, close around the object, lift it, move above the destination, and open to release. " + "the finger midpoint. Predict nearby waypoints using small pose changes. " "Follow the response schema exactly, with no explanatory text." ) diff --git a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py index 66a19b2ae7..a9af9b8a1e 100644 --- a/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py +++ b/isaaclab_arena_vlm_agent/tests/test_vlm_policy_configuration.py @@ -129,14 +129,14 @@ def test_configured_execution_and_partial_reset(monkeypatch, policy_type, cfg_ty @pytest.mark.parametrize("policy_type,cfg_type", POLICIES) -def test_injected_backend_remains_caller_owned(monkeypatch, policy_type, cfg_type): +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_not_called() + backend.close.assert_called_once() @pytest.mark.parametrize("policy_type,cfg_type", POLICIES) diff --git a/isaaclab_arena_vlm_agent/vlm_agent_policy.py b/isaaclab_arena_vlm_agent/vlm_agent_policy.py index c00fc40d71..6f545c180f 100644 --- a/isaaclab_arena_vlm_agent/vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/vlm_agent_policy.py @@ -181,7 +181,7 @@ def __init__( Args: config: Prompt, image encoding, and action-validation settings. - inference_backend: Caller-owned backend, or None to create a client owned by this policy. + 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. """ @@ -200,7 +200,6 @@ def __init__( action_adapter ) self.task_description = None - self._owns_inference_backend = inference_backend is None # Construct the client last so invalid adapters or schemas cannot leak a connection. self._inference_backend = ( inference_backend @@ -216,9 +215,8 @@ def __init__( ) def close(self) -> None: - """Close the inference backend when this policy created it.""" - if self._owns_inference_backend: - self._inference_backend.close() + """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. From e49c09ab7427b3ae9110e6df0975ea91770b876c Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 15:05:05 -0700 Subject: [PATCH 09/11] Move pick-and-place guidance into experiment config Signed-off-by: Xinjie Yao --- .../experiment_configs/droid_pnp_vlm_experiment.yaml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml index 3940b27bc7..72e0f2cffe 100644 --- a/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml +++ b/isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml @@ -20,7 +20,10 @@ runs: hdr: home_office_robolab environment_builder: &builder num_envs: 1 - language_instruction: Pick up the Rubik's cube and place it in the bowl. + 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 From f4a846e3101d58e7d8b5290a7dd0e84a280ea095 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 15:24:51 -0700 Subject: [PATCH 10/11] Extract VLM camera geometry helper Signed-off-by: Xinjie Yao --- .../observation_adapter.py | 108 ++++++++++-------- 1 file changed, 62 insertions(+), 46 deletions(-) diff --git a/isaaclab_arena_vlm_agent/observation_adapter.py b/isaaclab_arena_vlm_agent/observation_adapter.py index 0589520f35..45825efc93 100644 --- a/isaaclab_arena_vlm_agent/observation_adapter.py +++ b/isaaclab_arena_vlm_agent/observation_adapter.py @@ -104,52 +104,7 @@ def extract_vlm_input_from_observation(self, env, observation, env_id: int): ) 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) - 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 images, {"camera_geometry": geometry} + return images, {"camera_geometry": _extract_camera_geometry(env, images, env_id)} class ImageStateObservationAdapter(VLMObservationAdapter): @@ -164,3 +119,64 @@ def extract_vlm_input_from_observation(self, env, observation, env_id: int): _, state = self.state_adapter.extract_vlm_input_from_observation(env, observation, env_id) 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 From 35824a827bf1321f07056826e851f07ce5b3c126 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Mon, 26 Oct 2026 15:43:38 -0700 Subject: [PATCH 11/11] Extract camera calibration resize helper Signed-off-by: Xinjie Yao --- isaaclab_arena_vlm_agent/image_utils.py | 17 ++++++++++++++++- isaaclab_arena_vlm_agent/vlm_agent_policy.py | 11 ++++------- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/isaaclab_arena_vlm_agent/image_utils.py b/isaaclab_arena_vlm_agent/image_utils.py index 51d015be18..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 @@ -52,3 +52,18 @@ def image_to_jpeg_data_url(image: np.ndarray, max_width_or_height: int) -> tuple buffer = io.BytesIO() frame.save(buffer, format="JPEG", quality=85) 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/vlm_agent_policy.py b/isaaclab_arena_vlm_agent/vlm_agent_policy.py index 6f545c180f..ce64f428a5 100644 --- a/isaaclab_arena_vlm_agent/vlm_agent_policy.py +++ b/isaaclab_arena_vlm_agent/vlm_agent_policy.py @@ -24,7 +24,7 @@ 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.image_utils import image_to_jpeg_data_url, resize_camera_geometry from isaaclab_arena_vlm_agent.observation_adapter import ( ImageObservationAdapter, ImageStateObservationAdapter, @@ -256,12 +256,9 @@ def _request_content(self, env: gym.Env, observation: dict[str, Any], env_id: in 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", {}): - geometry = prompt_values["camera_geometry"][camera_name] - intrinsics = np.array(geometry["intrinsic_matrix"], dtype=float) - intrinsics[0] *= size[0] / image.shape[1] - intrinsics[1] *= size[1] / image.shape[0] - geometry["intrinsic_matrix"] = intrinsics.tolist() - geometry["image_size_px"] = list(size) + 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)})