Add experiment-runner support for VLMStepPolicy, VLMChunkPolicy, and VLMGoalPolicy with DROID PnP example

isaaclab_arena/embodiments/droid/droid.py

@@ -291,6 +291,27 @@ def _configure_newton_diff_ik(self) -> None:
         )
 
 
+@register_asset
+class DroidAbsoluteDifferentialIKEmbodiment(DroidDifferentialIKEmbodiment):
+    """DROID differential IK accepting unscaled root-relative XYZ/XYZW pose targets."""
+
+    registry_name = "droid_abs_ik"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._set_absolute_pose_actions()
+
+    def _configure_physics_backend(self, backend: PhysicsBackend) -> None:
+        """Apply backend IK settings, then accept unscaled absolute pose targets."""
+        super()._configure_physics_backend(backend)
+        self._set_absolute_pose_actions()
+
+    def _set_absolute_pose_actions(self) -> None:
+        """Keep the stored action config in absolute pose mode."""
+        self.action_config.arm_action.controller.use_relative_mode = False
+        self.action_config.arm_action.scale = 1.0
+
+
 @register_asset
 class DroidRelativeJointPositionEmbodiment(DroidEmbodimentBase):
     """Embodiment for the DROID setup with relative joint position action controller."""
@@ -569,6 +590,7 @@ class DroidCameraCfg(ArenaCameraCfg):
     """Configuration for cameras. DROID cameras are mounted with pre-set poses."""
 
     external_camera: CameraCfg = CameraCfg(
+        update_latest_camera_pose=True,
         prim_path="{ENV_REGEX_NS}/Robot/panda_link0/external_camera",
         height=720,
         width=1280,
@@ -582,6 +604,7 @@ class DroidCameraCfg(ArenaCameraCfg):
         offset=CameraCfg.OffsetCfg(pos=(0.05, 0.57, 0.66), rot=(-0.195, 0.399, 0.805, -0.393), convention="opengl"),
     )
     external_camera_2: CameraCfg = CameraCfg(
+        update_latest_camera_pose=True,
         prim_path="{ENV_REGEX_NS}/Robot/panda_link0/external_camera_2",
         height=720,
         width=1280,
@@ -595,6 +618,7 @@ class DroidCameraCfg(ArenaCameraCfg):
         offset=CameraCfg.OffsetCfg(pos=(0.05, -0.57, 0.66), rot=(0.399, -0.195, -0.393, 0.805), convention="opengl"),
     )
     wrist_camera: CameraCfg = CameraCfg(
+        update_latest_camera_pose=True,
         prim_path="{ENV_REGEX_NS}/Robot/Gripper/Robotiq_2F_85/base_link/wrist_camera",
         height=720,
         width=1280,

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, :]

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,16 +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()
     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)

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.
 

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()

isaaclab_arena_environments/experiment_configs/droid_pnp_vlm_experiment.yaml

@@ -0,0 +1,55 @@
+# 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
+      model: openai/openai/gpt-6-astra
+      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
+      model: openai/openai/gpt-6-astra
+      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
+      model: openai/openai/gpt-6-astra
+      vlm_embodiment_adapter: droid
+      observation_modality: image_state
+      max_steps_per_goal: 30
+    rollout_limit: *rollout

isaaclab_arena_vlm_agent/absolute_action_adapter.py

@@ -7,54 +7,62 @@
 
 import numpy as np
 
-from isaaclab_arena_vlm_agent.action_adapter import ChunkActionAdapter, GoalActionAdapter
+from isaaclab_arena_vlm_agent.action_adapter import AgentActionAdapter, ChunkActionAdapter, GoalActionAdapter
 from isaaclab_arena_vlm_agent.command_indices import AbsoluteCommandIndices, read_absolute_command_indices
 from isaaclab_arena_vlm_agent.commands import Command, CommandChunk
 from isaaclab_arena_vlm_agent.motion_generator import MotionGenerator
 from isaaclab_arena_vlm_agent.response_errors import InvalidAgentResponseError
 
 
-class AbsoluteChunkActionAdapter(ChunkActionAdapter):
-    """Decode absolute command chunks; use chunk_size=1 for one prediction per step."""
+class _AbsoluteCommandAdapter(AgentActionAdapter):
+    """Decode a single absolute command in native controller order."""
 
-    response_model = CommandChunk
+    response_model = Command
 
-    def __init__(self, chunk_size: int):
-        super().__init__(chunk_size)
+    def __init__(self):
         self._command_indices = None
 
     def is_robot_action_space_supported(self, env) -> bool:
-        """Resolve command indices from the active controller before policy inference."""
+        """Read absolute command indices before policy inference."""
         self._command_indices = read_absolute_command_indices(env)
         return self._command_indices is not None
 
+    def decode_command(self, response: Command, arena_world, env_id: int) -> np.ndarray:
+        """Decode and validate one absolute target."""
+        return _decode_absolute_command(response, self._command_indices)
+
+
+class AbsoluteStepActionAdapter(_AbsoluteCommandAdapter):
+    """Decode a single absolute command in native controller order."""
+
+
+class AbsoluteChunkActionAdapter(_AbsoluteCommandAdapter, ChunkActionAdapter):
+    """Decode absolute command chunks in native controller order."""
+
+    response_model = CommandChunk
+
+    def __init__(self, chunk_size: int):
+        _AbsoluteCommandAdapter.__init__(self)
+        ChunkActionAdapter.__init__(self, chunk_size)
+
     def decode_command(self, response: CommandChunk, arena_world, env_id: int) -> np.ndarray:
         """Validate and stack absolute commands in execution order."""
         if len(response.commands) != self.chunk_size:
             raise InvalidAgentResponseError(f"Expected {self.chunk_size} commands, got {len(response.commands)}")
         commands = []
         for command in response.commands:
-            commands.append(_decode_absolute_command(command, self._command_indices))
+            commands.append(super().decode_command(command, arena_world, env_id))
         return np.stack(commands)
 
 
-class AbsoluteGoalActionAdapter(GoalActionAdapter):
+class AbsoluteGoalActionAdapter(_AbsoluteCommandAdapter, GoalActionAdapter):
     """Decode absolute targets for an explicitly supplied motion generator."""
 
     response_model = Command
 
     def __init__(self, motion_generator: MotionGenerator):
-        super().__init__(motion_generator)
-        self._command_indices = None
-
-    def is_robot_action_space_supported(self, env) -> bool:
-        """Resolve command indices from the active controller before policy inference."""
-        self._command_indices = read_absolute_command_indices(env)
-        return self._command_indices is not None
-
-    def decode_command(self, response: Command, arena_world, env_id: int) -> np.ndarray:
-        """Decode a validated response into an absolute target for the motion generator."""
-        return _decode_absolute_command(response, self._command_indices)
+        _AbsoluteCommandAdapter.__init__(self)
+        GoalActionAdapter.__init__(self, motion_generator)
 
 
 def _decode_absolute_command(response: Command, indices: AbsoluteCommandIndices | None) -> np.ndarray:

isaaclab_arena_vlm_agent/droid_adapter.py

@@ -0,0 +1,98 @@
+# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""DROID absolute-pose action, motion-generation, and prompt defaults."""
+
+from isaaclab_arena_vlm_agent.absolute_action_adapter import (
+    AbsoluteChunkActionAdapter,
+    AbsoluteGoalActionAdapter,
+    AbsoluteStepActionAdapter,
+)
+from isaaclab_arena_vlm_agent.action_adapter import AgentActionAdapter
+from isaaclab_arena_vlm_agent.command_indices import AbsoluteCommandIndices, GripperCommandIndex, PoseCommandIndices
+from isaaclab_arena_vlm_agent.pose_target_motion_generator import PoseTargetMotionGenerator
+
+
+class _DroidPoseActionAdapter:
+    """Require the native pose and gripper layout described by the DROID prompt."""
+
+    def is_robot_action_space_supported(self, env) -> bool:
+        return super().is_robot_action_space_supported(env) and self._command_indices == AbsoluteCommandIndices(
+            action_dim=8,
+            poses=(PoseCommandIndices(slice(0, 7), "robot", "base_link"),),
+            joints=(),
+            grippers=(GripperCommandIndex(7, "gripper_action", 0.0, 1.0),),
+        )
+
+
+class _DroidStepActionAdapter(_DroidPoseActionAdapter, AbsoluteStepActionAdapter):
+    """Decode a DROID pose command."""
+
+
+class _DroidChunkActionAdapter(_DroidPoseActionAdapter, AbsoluteChunkActionAdapter):
+    """Decode a chunk of DROID pose commands."""
+
+
+class _DroidGoalActionAdapter(_DroidPoseActionAdapter, AbsoluteGoalActionAdapter):
+    """Decode a DROID pose goal."""
+
+
+class DroidVLMAdapter:
+    """Provide VLM defaults for the droid_abs_ik embodiment."""
+
+    def make_step_action_adapter(self) -> AbsoluteStepActionAdapter:
+        """Create a decoder for one absolute pose command."""
+        return _DroidStepActionAdapter()
+
+    def make_chunk_action_adapter(self, chunk_size: int) -> AbsoluteChunkActionAdapter:
+        """Create a decoder for the requested number of absolute pose commands."""
+        return _DroidChunkActionAdapter(chunk_size)
+
+    def make_goal_action_adapter(
+        self, position_tolerance_m: float, rotation_tolerance_rad: float
+    ) -> AbsoluteGoalActionAdapter:
+        """Create pose-goal execution with DROID's Robotiq gripper measurements."""
+        generator = PoseTargetMotionGenerator(
+            grippers={"gripper_action": _make_gripper()},
+            position_tolerance_m=position_tolerance_m,
+            rotation_tolerance_rad=rotation_tolerance_rad,
+        )
+        return _DroidGoalActionAdapter(generator)
+
+    def get_system_prompt(self, action_adapter: AgentActionAdapter) -> str:
+        """Describe the DROID pose contract and the selected response schema."""
+        if isinstance(action_adapter, AbsoluteChunkActionAdapter):
+            response = (
+                f"Return a JSON object with a commands list of exactly {action_adapter.chunk_size} commands, "
+                "executed one per simulation step. Each command has a values list of eight numbers: "
+            )
+        else:
+            assert isinstance(
+                action_adapter, (AbsoluteStepActionAdapter, AbsoluteGoalActionAdapter)
+            ), "Provide system_prompt when injecting a custom action adapter"
+            response = "Return one command as a JSON object with a values list of eight numbers: "
+        return (
+            "Control DROID to complete the supplied task using the current camera images and robot state. "
+            + response
+            + "absolute XYZ in meters and a unit XYZW quaternion for the gripper base_link body, expressed in "
+            "the robot root frame, followed by the gripper command (0 opens, 1 closes). Read the current "
+            "camera_geometry for each named image: image_size_px is [width, height], intrinsic_matrix is K "
+            "in the transmitted image's pixels, and camera poses map optical coordinates (+X right, +Y down, "
+            "+Z forward) into the named world or robot root frame. For a camera point p_C, "
+            "p_R = rotate(quaternion_root_xyzw, p_C) + position_root_m; pixels project as K p_C / p_C.z. "
+            "All quaternions in input and output use XYZW. policy.eef_pos is in world meters and "
+            "policy.eef_quat is the gripper's world orientation in XYZW. Read the current "
+            "commanded-body pose from command_body_poses.arm_action when state is provided. This body is not "
+            "the finger midpoint. Predict nearby waypoints using small pose changes. 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()