Add WBC-AGILE end-to-end velocity policy for G1

isaaclab_arena/embodiments/g1/g1.py

@@ -136,6 +136,32 @@ def __init__(
         self.camera_config._camera_offset = camera_offset
 
 
+@register_asset
+class G1WBCAgileJointEmbodiment(G1EmbodimentBase):
+    """Embodiment for the G1 robot with AGILE WBC policy and direct joint upperbody control.
+
+    By default uses tiled camera for efficient parallel evaluation.
+    """
+
+    name = "g1_wbc_agile_joint"
+
+    def __init__(
+        self,
+        enable_cameras: bool = False,
+        initial_pose: Pose | None = None,
+        camera_offset: Pose | None = _DEFAULT_G1_CAMERA_OFFSET,
+        use_tiled_camera: bool = True,
+    ):
+        super().__init__(enable_cameras, initial_pose)
+        self.action_config = G1WBCAgileJointActionCfg()
+        self.observation_config = G1WBCJointObservationsCfg()
+        self.observation_config.policy.concatenate_terms = self.concatenate_observation_terms
+        self.observation_config.wbc.concatenate_terms = self.concatenate_observation_terms
+        self.event_config = G1WBCJointEventCfg()
+        self.camera_config._is_tiled_camera = use_tiled_camera
+        self.camera_config._camera_offset = camera_offset
+
+
 @configclass
 class G1SceneCfg:
 
@@ -599,6 +625,13 @@ class G1WBCJointActionCfg:
     g1_action: ActionTermCfg = G1DecoupledWBCJointActionCfg(asset_name="robot", joint_names=[".*"])
 
 
+@configclass
+class G1WBCAgileJointActionCfg:
+    """Action specifications for the MDP, for G1 AGILE WBC action."""
+
+    g1_action: ActionTermCfg = G1DecoupledWBCJointActionCfg(asset_name="robot", joint_names=[".*"], wbc_version="agile")
+
+
 @configclass
 class G1WBCPinkActionCfg:
     """Action specifications for the MDP, for G1 WBC action."""

isaaclab_arena_g1/g1_env/mdp/actions/g1_decoupled_wbc_joint_action.py

@@ -15,7 +15,7 @@
 from isaaclab.assets.articulation import Articulation
 from isaaclab.managers.action_manager import ActionTerm
 
-from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.config.configs import HomieV2Config
+from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.config.configs import AgileConfig, HomieV2Config
 from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.policy.policy_constants import (
     G1_NUM_JOINTS,
     NUM_BASE_HEIGHT_CMD,
@@ -70,6 +70,8 @@ def __init__(self, cfg: G1DecoupledWBCJointActionCfg, env: ManagerBasedEnv):
 
         if self._wbc_version == "homie_v2":
             wbc_config = HomieV2Config()
+        elif self._wbc_version == "agile":
+            wbc_config = AgileConfig()
         else:
             raise ValueError(f"Invalid WBC version: {self._wbc_version}")
 

isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/config/configs.py

@@ -71,3 +71,24 @@ class HomieV2Config(BaseConfig):
 
     policy_config_path: str = "config/g1_homie_v2.yaml"
     """Policy related configuration to specify inputs/outputs dim"""
+
+
+@dataclass
+class AgileConfig(BaseConfig):
+    """Config for the WBC-AGILE end-to-end velocity policy for G1."""
+
+    # WBC Configuration
+    wbc_version: Literal["agile"] = "agile"
+    """Version of the whole body controller."""
+
+    wbc_model_path: str = (
+        "https://github.com/nvidia-isaac/WBC-AGILE/raw/v1.2/agile/data/policy/velocity_g1/unitree_g1_velocity_e2e.onnx"
+    )
+    """Path to WBC model file (GitHub URL, resolved via retrieve_file_path)"""
+
+    # Robot Configuration
+    enable_waist: bool = False
+    """Whether to include waist joints in IK."""
+
+    policy_config_path: str = "config/g1_agile.yaml"
+    """Policy related configuration to specify inputs/outputs dim"""

isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/config/g1_agile.yaml

@@ -0,0 +1,61 @@
+# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Joint ordering for ONNX model inputs (29 body joints in agile order).
+# Must match the element_names from unitree_g1_velocity_e2e.yaml.
+onnx_input_joint_names:
+  - left_hip_pitch_joint
+  - right_hip_pitch_joint
+  - waist_yaw_joint
+  - left_hip_roll_joint
+  - right_hip_roll_joint
+  - waist_roll_joint
+  - left_hip_yaw_joint
+  - right_hip_yaw_joint
+  - waist_pitch_joint
+  - left_knee_joint
+  - right_knee_joint
+  - left_shoulder_pitch_joint
+  - right_shoulder_pitch_joint
+  - left_ankle_pitch_joint
+  - right_ankle_pitch_joint
+  - left_shoulder_roll_joint
+  - right_shoulder_roll_joint
+  - left_ankle_roll_joint
+  - right_ankle_roll_joint
+  - left_shoulder_yaw_joint
+  - right_shoulder_yaw_joint
+  - left_elbow_joint
+  - right_elbow_joint
+  - left_wrist_roll_joint
+  - right_wrist_roll_joint
+  - left_wrist_pitch_joint
+  - right_wrist_pitch_joint
+  - left_wrist_yaw_joint
+  - right_wrist_yaw_joint
+
+# Joint ordering for ONNX model outputs (14 controlled joints in agile output order).
+# Must match the action_joint_pos element_names from unitree_g1_velocity_e2e.yaml.
+controlled_joint_names:
+  - left_hip_pitch_joint
+  - right_hip_pitch_joint
+  - left_hip_roll_joint
+  - right_hip_roll_joint
+  - waist_roll_joint
+  - left_hip_yaw_joint
+  - right_hip_yaw_joint
+  - waist_pitch_joint
+  - left_knee_joint
+  - right_knee_joint
+  - left_ankle_pitch_joint
+  - right_ankle_pitch_joint
+  - left_ankle_roll_joint
+  - right_ankle_roll_joint
+
+num_actions: 14
+num_body_joints: 29
+
+# Initial commands
+cmd_init: [0.0, 0.0, 0.0]

isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/policy/g1_agile_policy.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pathlib
+import torch
+from typing import Any
+
+from isaaclab.utils.assets import retrieve_file_path
+
+from isaaclab_arena_g1.g1_env.robot_model import RobotModel
+from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.policy.base import WBCPolicy
+from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.utils.homie_utils import load_config
+from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.utils.onnx_utils import OnnxInferenceSession
+
+# ONNX feedback state keys and their per-environment shapes (excluding batch dim).
+_STATE_KEYS_AND_SHAPES: dict[str, tuple[int, ...]] = {
+    "last_actions": (14,),
+    "base_ang_vel_history": (5, 3),
+    "projected_gravity_history": (5, 3),
+    "velocity_commands_history": (5, 3),
+    "controlled_joint_pos_history": (5, 14),
+    "controlled_joint_vel_history": (5, 14),
+    "actions_history": (5, 14),
+}
+
+
+class G1AgilePolicy(WBCPolicy):
+    """G1 robot policy using the WBC-AGILE end-to-end neural network.
+
+    This policy uses a single ONNX model that takes raw sensor inputs and
+    manages observation history internally via feedback connections. The model
+    outputs target joint positions along with per-joint Kp/Kd gains for 14
+    controlled joints (legs + waist_roll + waist_pitch).
+    """
+
+    def __init__(self, robot_model: RobotModel, config_path: str, model_path: str, num_envs: int = 1):
+        """Initialize G1AgilePolicy.
+
+        Args:
+            robot_model: Robot model containing joint ordering info.
+            config_path: Path to policy YAML configuration file (relative to wbc_policy dir).
+            model_path: Path to the ONNX model file. Can be an S3/Nucleus URL
+                (resolved and cached by retrieve_file_path) or a local path.
+            num_envs: Number of environments.
+        """
+        parent_dir = pathlib.Path(__file__).parent.parent
+        self.config = load_config(str(parent_dir / config_path))
+        self.robot_model = robot_model
+        self.num_envs = num_envs
+
+        # Resolve model path via OV asset API (handles S3 download + local caching).
+        # Same pattern as G1HomiePolicyV2: retrieve_file_path returns a local path
+        # (absolute for S3 cache, relative for local files), then join with parent_dir.
+        model_local_path = retrieve_file_path(model_path, force_download=True)
+        model_full_path = parent_dir / model_local_path
+        self.session = OnnxInferenceSession(str(model_full_path))
+
+        # Build joint index mappings between WBC order and agile ONNX order
+        self._build_joint_mappings()
+
+        # Initialize state
+        self._init_state()
+
+    def _build_joint_mappings(self):
+        """Build index mappings between WBC joint order and agile ONNX joint order."""
+        wbc_order = self.robot_model.wbc_g1_joints_order  # {joint_name: wbc_index}
+
+        # Mapping for ONNX input: indices into the WBC-ordered observation to select
+        # the 29 body joints in the order the ONNX model expects.
+        onnx_input_names = self.config["onnx_input_joint_names"]
+        self.wbc_to_agile_input = np.array([wbc_order[name] for name in onnx_input_names])
+
+        # Mapping for ONNX output: for each of the 14 agile output joints, the
+        # position in the 15-element lower_body array to write to.
+        controlled_names = self.config["controlled_joint_names"]
+        lower_body_indices = self.robot_model.get_joint_group_indices("lower_body")
+        self.agile_idx_to_lower_body_idx = np.array(
+            [lower_body_indices.index(wbc_order[name]) for name in controlled_names]
+        )
+
+        self.num_lower_body = len(lower_body_indices)
+
+    def _init_state(self):
+        """Initialize all per-environment state variables."""
+        self.observation = None
+        self.cmd = np.tile(self.config["cmd_init"], (self.num_envs, 1))
+
+        # Batched ONNX feedback state: each array has shape (num_envs, ...).
+        self.state = {
+            key: np.zeros((self.num_envs, *shape), dtype=np.float32) for key, shape in _STATE_KEYS_AND_SHAPES.items()
+        }
+
+    def reset(self, env_ids: torch.Tensor):
+        """Reset the policy state for the given environment ids.
+
+        Args:
+            env_ids: The environment ids to reset.
+        """
+        idx = env_ids.cpu().numpy()
+        for key, shape in _STATE_KEYS_AND_SHAPES.items():
+            self.state[key][idx] = np.zeros(shape, dtype=np.float32)
+        self.cmd[idx] = self.config["cmd_init"]
+
+    def set_observation(self, observation: dict[str, Any]):
+        """Store the current observation for the next get_action call.
+
+        Args:
+            observation: Dictionary containing robot state from prepare_observations().
+        """
+        self.observation = observation
+
+    def set_goal(self, goal: dict[str, Any]):
+        """Set the goal for the policy.
+
+        Args:
+            goal: Dictionary containing goals. Supported keys:
+                - "navigate_cmd": velocity command array of shape (num_envs, 3)
+        """
+        if "navigate_cmd" in goal:
+            self.cmd = goal["navigate_cmd"]
+
+    def get_action(self, time: float | None = None) -> dict[str, Any]:
+        """Compute and return the next action based on current observation.
+
+        Returns:
+            Dictionary with "body_action" key containing joint position targets
+            of shape (num_envs, num_lower_body_joints).
+        """
+        if self.observation is None:
+            raise ValueError("No observation set. Call set_observation() first.")
+
+        obs = self.observation
+
+        # Build batched ONNX inputs (all envs at once)
+        ort_inputs = {
+            "root_link_quat_w": obs["floating_base_pose"][:, 3:7].astype(np.float32),
+            "root_ang_vel_b": obs["floating_base_vel"][:, 3:6].astype(np.float32),
+            "velocity_commands": self.cmd.astype(np.float32),
+            "joint_pos": obs["q"][:, self.wbc_to_agile_input].astype(np.float32),
+            "joint_vel": obs["dq"][:, self.wbc_to_agile_input].astype(np.float32),
+            **{key: self.state[key] for key in _STATE_KEYS_AND_SHAPES},
+        }
+
+        # Run batched inference
+        result = self.session.run(ort_inputs)
+
+        # Update feedback state for next step
+        for key in _STATE_KEYS_AND_SHAPES:
+            self.state[key] = result[f"{key}_out"]
+
+        # Map 14 agile output joints to the 15-joint lower_body array.
+        # waist_yaw (not controlled by agile) stays at 0.0.
+        body_action = np.zeros((self.num_envs, self.num_lower_body), dtype=np.float32)
+        body_action[:, self.agile_idx_to_lower_body_idx] = result["action_joint_pos"]
+
+        return {"body_action": body_action}

isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/policy/g1_homie_policy.py

@@ -7,14 +7,13 @@
 import numpy as np
 import pathlib
 import torch
-from collections.abc import Callable
 from typing import Any
 
-import onnxruntime as ort
 from isaaclab.utils.assets import retrieve_file_path
 
 from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.policy.base import WBCPolicy
 from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.utils.homie_utils import get_gravity_orientation, load_config
+from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.utils.onnx_utils import OnnxInferenceSession
 
 
 class G1HomiePolicyV2(WBCPolicy):
@@ -38,8 +37,8 @@ def __init__(self, robot_model, config_path: str, model_path: str, num_envs: int
         model_path_1_local = retrieve_file_path(model_path_1, force_download=True)
         model_path_2_local = retrieve_file_path(model_path_2, force_download=True)
 
-        self.policy_1 = self.load_onnx_policy(str(parent_dir / model_path_1_local))
-        self.policy_2 = self.load_onnx_policy(str(parent_dir / model_path_2_local))
+        self.policy_1 = OnnxInferenceSession(str(parent_dir / model_path_1_local))
+        self.policy_2 = OnnxInferenceSession(str(parent_dir / model_path_2_local))
 
         # Initialize observation history buffer
         self.observation = None
@@ -81,26 +80,6 @@ def reset(self, env_ids: torch.Tensor):
         self.pitch_cmd = self.config["rpy_cmd"][1]
         self.yaw_cmd = self.config["rpy_cmd"][2]
 
-    def load_onnx_policy(self, model_path: str) -> Callable[[torch.Tensor], torch.Tensor]:
-        """Load the ONNX policy from the model path.
-
-        Args:
-            model_path: The path to the ONNX policy model
-
-        Returns:
-            The ONNX policy model runnable for forward pass.
-        """
-        model = ort.InferenceSession(model_path)
-
-        def run_inference(input_tensor):
-            ort_inputs = {model.get_inputs()[0].name: input_tensor.cpu().numpy()}
-            ort_outs = model.run(None, ort_inputs)
-            return torch.tensor(ort_outs[0], device="cpu")
-
-        print(f"Successfully loaded ONNX policy from {model_path}")
-
-        return run_inference
-
     def compute_observation(self, observation: dict[str, Any]) -> tuple[np.ndarray, int]:
         """Compute the observation vector from current state"""
         # Get body joint indices (excluding waist roll and pitch)
@@ -245,7 +224,9 @@ def get_action(self) -> dict[str, Any]:
                 # Use walking policy for movement commands
                 policy = self.policy_2
 
-            self.action = policy(self.obs_tensor).detach().numpy()
+            ort_inputs = {policy.input_names[0]: self.obs_tensor.cpu().numpy()}
+            result = policy.run(ort_inputs)
+            self.action = list(result.values())[0]
 
         # Transform action to target_dof_pos
         assert self.use_policy_action