project init

.gitignore

@@ -0,0 +1,9 @@
+.venv/
+*.vscode/
+build/
+*.egg-info/
+*__pycache__/
+*logs/
+*outputs/
+dataset/
+models/

MANIFEST.in

@@ -0,0 +1 @@
+recursive-include source/engineai_lab/assets *

README.md

@@ -0,0 +1,153 @@
+# EngineAI-Lab
+
+**EngineAI Lab** is a python package for training and deploying policies for EngineAI Robots using Isaac Lab and Isaac Sim.
+
+|Training| Sim2Sim |Deploy|
+|--------|--------|--------|
+![training](./docs/training.gif)|![sim2sim](./docs/sim2sim.gif)|<img src="./docs/deploy.gif" height="180"/>
+
+# Structure
+
+```
+engineai-lab
+├── config
+├── dataset
+│   ├── config
+│   └── data
+├── scripts
+└── source
+    └── engineai_lab
+        ├── algorithms      
+        ├── assets
+        │   └── pm01
+        │       ├── meshes
+        │       └── urdf
+        ├── robots
+        ├── tasks
+        │   └── velocity
+        │       ├── config
+        │       │   └── pm01
+        │       └── mdp
+        └── utils
+```
+
+## QUICKSTART
+
+### 1. Create a Conda Environment
+
+Create and activate a new environment with Python 3.11:
+
+  ```bash
+  conda create -n engineai_lab python=3.11
+  conda activate engineai_lab
+  ```
+
+### 2. Install Prerequisites
+
+- **Install Isaac Sim**
+
+  Follow the official installation guide: [Isaac Lab - Pip Installation](https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/pip_installation.html#installing-dependencies).
+
+    *Since you've already created the engineai_lab environment, follow the guide from "Installing Dependencies" up to (but not including) the "Installing Isaac Lab" section.*
+
+- **Clone & Setup Isaac Lab**
+
+  Clone the repository and switch to the recommended branch:
+
+  ```bash
+  git clone https://github.com/isaac-sim/IsaacLab.git
+  cd IsaacLab
+  git checkout 4df6560e
+  ./isaaclab -i rsl_rl   # Install rsl-rl dependency
+  ```
+
+We highly recommend using the main branch`(4df6560e)` of Isaac Lab, as it can support rsl-rl-lib >= 5.0 and Isaac Sim >= 5.0 .
+
+### 3. Install this Package
+
+Once the prerequisites are set up, install the package in editable mode:
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+### Supported Robots
+
+This repository currently supports the following environments from the EngineAI Robots family:
+
+|Robot| Task |Description|
+|--------|--------|--------|
+PM01|`Flat-PM01-v0`|Basic flat-terrain locomotion
+PM01|`Flat-AMP-PM01-v0`|AMP-based motion imitation on flat terrain
+
+*More robots and environments are coming soon!*
+
+### Training a Policy
+
+```
+ python scripts/train.py --task=Flat-PM01-v0  --num_envs 4096 --headless --run_name <name>  
+ python scripts/play.py  --task=Flat-PM01-v0 --num_envs 128 --load_run <name> 
+```
+
+### Evaluating a Policy
+
+```
+ python scripts/train.py --task=Flat-AMP-PM01-v0  --num_envs 4096 --headless --run_name <name> 
+ python scripts/play.py --task=Flat-AMP-PM01-v0 --num_envs 128 --load_run <name>
+```
+
+Replace `<name>` with the name of your training run (found in logs/rsl_rl/).
+
+### Deployment
+
+To deploy a trained policy on real hardware, convert it to the MNN format for efficient inference.
+
+#### 1. Export PyTorch Policy to ONNX
+
+(Ensure your training script supports ONNX export)
+
+#### 2. Build [MNN-Converter](https://mnn-docs.readthedocs.io/en/latest/start/quickstart_cpp.html?highlight=mnn+converter)
+
+```bash
+git clone https://github.com/alibaba/mnn
+cd mnn
+mkdir build && cd build
+cmake .. -DMNN_BUILD_CONVERTER=ON
+make -j8
+```
+
+#### 3. Convert ONNX into MNN
+
+```bash
+./MNNConvert -f ONNX \
+  --modelFile path_to_your_policy.onnx \
+  --MNNModel your_policy.mnn \
+  --bizCode MNN
+```
+
+For detailed instructions on integrating the MNN model with EngineAI robots, see:
+[engineai_robotics_native_sdk](https://github.com/engineai-robotics/engineai_robotics_native_sdk).
+
+## Support
+
+If you have any questions about using this repository, we're here to help!
+
+- **Report Issues**: Found a bug or have a feature request. Please open a new issue on our [GitHub Issues](https://github.com/engineai-robotics/engineai_lab/issues) page.
+- **Email Us**: For general inquiries or collaboration opportunities, feel free to reach out at [info@engineai.com.cn](mailto:info@engineai.com.cn).
+
+## License
+
+EngineAI-Lab is released under [BSD-3 License](LICENSE).
+
+## Acknowledgement
+
+This repository is built upon the support and contributions of the following open-source projects. Special thanks to:
+
+- [**IsaacLab**](https://github.com/isaac-sim/IsaacLab) — The foundational framework for training and running simulation experiments.
+- [**rsl_rl**](https://github.com/leggedrobotics/rsl_rl) — High-performance reinforcement learning library for legged robots.
+- [**AMP_for_hardware**](https://github.com/escontra/AMP_for_hardware) — Implementation of Adversarial Motion Priors (AMP) for sim-to-real transfer.
+- [**BeyondMimic**](https://github.com/HybridRobotics/whole_body_tracking) — Inspiration for project structure and valuable feature implementations.
+- [**MNN**](https://github.com/alibaba/mnn) — Lightweight, high-performance inference engine for on-device deployment.
+- [**engineai_robotics_native_sdk**](https://github.com/engineai-robotics/engineai_robotics_native_sdk) — Official SDK for deploying policies on EngineAI robotic hardware.

config/extension.toml

@@ -0,0 +1,35 @@
+[package]
+
+# Semantic Versioning is used: https://semver.org/
+version = "0.1.0"
+
+# Description
+category = "isaaclab"
+readme  = "README.md"
+
+title = "Extension Template"
+author = "Isaac Lab Project Developers"
+maintainer = "Isaac Lab Project Developers"
+description="Extension Template for Isaac Lab"
+repository = "https://github.com/isaac-sim/IsaacLabExtensionTemplate.git"
+keywords = ["extension", "template", "isaaclab"]
+
+[dependencies]
+"isaaclab" = {}
+"isaaclab_assets" = {}
+"isaaclab_mimic" = {}
+"isaaclab_rl" = {}
+"isaaclab_tasks" = {}
+# NOTE: Add additional dependencies here
+
+[[python.module]]
+name = "engineai_lab"
+
+[isaaclab_settings]
+  # TODO: Uncomment and list any apt dependencies here.
+#       If none, leave it commented out.
+# apt_deps = ["example_package"]
+# TODO: Uncomment and provide path to a ros_ws
+#       with rosdeps to be installed. If none,
+#       leave it commented out.
+# ros_ws = "path/from/extension_root/to/ros_ws"

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "wheel", "toml"]
+build-backend = "setuptools.build_meta"

scripts/cli_args.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import argparse
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg
+
+
+def add_rsl_rl_args(parser: argparse.ArgumentParser):
+    """Add RSL-RL arguments to the parser.
+
+    Args:
+        parser: The parser to add the arguments to.
+    """
+    # create a new argument group
+    arg_group = parser.add_argument_group("rsl_rl", description="Arguments for RSL-RL agent.")
+    # -- experiment arguments
+    arg_group.add_argument(
+        "--experiment_name", type=str, default=None, help="Name of the experiment folder where logs will be stored."
+    )
+    arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.")
+    # -- load arguments
+    arg_group.add_argument("--resume", type=bool, default=None, help="Whether to resume from a checkpoint.")
+    arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.")
+    arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.")
+    # -- logger arguments
+    arg_group.add_argument(
+        "--logger", type=str, default=None, choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use."
+    )
+    arg_group.add_argument(
+        "--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune."
+    )
+    arg_group.add_argument(
+        "--wandb_path", type=str, default=None, help="Name of the logging project when using wandb or neptune."
+    )
+    arg_group.add_argument(
+        "--rl_device",
+        type=str,
+        default=None,
+        help="Device used by the RSL-RL policy/optimizer, e.g. cpu, cuda, or cuda:0.",
+    )
+
+
+def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPolicyRunnerCfg:
+    """Parse configuration for RSL-RL agent based on inputs.
+
+    Args:
+        task_name: The name of the environment.
+        args_cli: The command line arguments.
+
+    Returns:
+        The parsed configuration for RSL-RL agent based on inputs.
+    """
+    from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry
+
+    # load the default configuration
+    rslrl_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point")
+    rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli)
+    return rslrl_cfg
+
+
+def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Namespace):
+    """Update configuration for RSL-RL agent based on inputs.
+
+    Args:
+        agent_cfg: The configuration for RSL-RL agent.
+        args_cli: The command line arguments.
+
+    Returns:
+        The updated configuration for RSL-RL agent based on inputs.
+    """
+    # override the default configuration with CLI arguments
+    if hasattr(args_cli, "seed") and args_cli.seed is not None:
+        agent_cfg.seed = args_cli.seed
+    if args_cli.resume is not None:
+        agent_cfg.resume = args_cli.resume
+    if args_cli.load_run is not None:
+        agent_cfg.load_run = args_cli.load_run
+    if args_cli.checkpoint is not None:
+        agent_cfg.load_checkpoint = args_cli.checkpoint
+    if args_cli.run_name is not None:
+        agent_cfg.run_name = args_cli.run_name
+    if args_cli.logger is not None:
+        agent_cfg.logger = args_cli.logger
+    if getattr(args_cli, "rl_device", None) is not None:
+        agent_cfg.device = args_cli.rl_device
+    # set the project name for wandb and neptune
+    if agent_cfg.logger in {"wandb", "neptune"} and args_cli.log_project_name:
+        agent_cfg.wandb_project = args_cli.log_project_name
+        agent_cfg.neptune_project = args_cli.log_project_name
+
+    return agent_cfg

scripts/gen2_check_rl_readiness.py

@@ -0,0 +1,326 @@
+#!/usr/bin/env python3
+"""Static readiness checks for the Gen2 IsaacLab locomotion asset."""
+
+from __future__ import annotations
+
+import ast
+import math
+import re
+import xml.etree.ElementTree as ET
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+URDF_PATH = REPO_ROOT / "source" / "gen2_lab" / "assets" / "robot_simplified_collision.urdf"
+GEN2_ROBOT_CFG = REPO_ROOT / "source" / "engineai_lab" / "robots" / "gen2.py"
+GEN2_ENV_CFG = REPO_ROOT / "source" / "engineai_lab" / "tasks" / "velocity" / "config" / "gen2" / "flat_env_cfg.py"
+FOOT_LINKS = ("left_leg_link_6", "right_leg_link_6")
+
+
+@dataclass
+class Box:
+    name: str
+    center: np.ndarray
+    axes: np.ndarray
+    half: np.ndarray
+    vertices: np.ndarray
+
+
+def _parse_floats(text: str | None, default: tuple[float, ...]) -> tuple[float, ...]:
+    if not text:
+        return default
+    return tuple(float(value) for value in text.split())
+
+
+def _rpy_matrix(rpy: tuple[float, float, float]) -> np.ndarray:
+    roll, pitch, yaw = rpy
+    cr, sr = math.cos(roll), math.sin(roll)
+    cp, sp = math.cos(pitch), math.sin(pitch)
+    cy, sy = math.cos(yaw), math.sin(yaw)
+    return np.array(
+        [
+            [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
+            [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
+            [-sp, cp * sr, cp * cr],
+        ],
+        dtype=float,
+    )
+
+
+def _axis_angle(axis: tuple[float, float, float], angle: float) -> np.ndarray:
+    vec = np.asarray(axis, dtype=float)
+    norm = np.linalg.norm(vec)
+    if norm < 1.0e-12:
+        return np.eye(3)
+    x, y, z = vec / norm
+    c, s = math.cos(angle), math.sin(angle)
+    one_c = 1.0 - c
+    return np.array(
+        [
+            [c + x * x * one_c, x * y * one_c - z * s, x * z * one_c + y * s],
+            [y * x * one_c + z * s, c + y * y * one_c, y * z * one_c - x * s],
+            [z * x * one_c - y * s, z * y * one_c + x * s, c + z * z * one_c],
+        ],
+        dtype=float,
+    )
+
+
+def _transform(xyz=(0.0, 0.0, 0.0), rpy=(0.0, 0.0, 0.0)) -> np.ndarray:
+    matrix = np.eye(4)
+    matrix[:3, :3] = _rpy_matrix(rpy)
+    matrix[:3, 3] = np.asarray(xyz, dtype=float)
+    return matrix
+
+
+def _load_default_joint_positions() -> dict[str, float]:
+    text = GEN2_ROBOT_CFG.read_text(encoding="utf-8")
+    match = re.search(r"joint_pos=\{(.*?)\},\n\s*joint_vel=", text, flags=re.S)
+    if match is None:
+        return {}
+    return ast.literal_eval("{" + match.group(1) + "}")
+
+
+def _load_base_init_z() -> float | None:
+    text = GEN2_ROBOT_CFG.read_text(encoding="utf-8")
+    match = re.search(r"pos=\(([^)]*)\)", text)
+    if match is None:
+        return None
+    return float(match.group(1).split(",")[2].strip())
+
+
+def _load_base_height_target() -> float | None:
+    if not GEN2_ENV_CFG.exists():
+        return None
+    text = GEN2_ENV_CFG.read_text(encoding="utf-8")
+    match = re.search(r"GEN2_BASE_HEIGHT_TARGET\s*=\s*([0-9.]+)", text)
+    if match is None:
+        return None
+    return float(match.group(1))
+
+
+def _link_tree(root: ET.Element):
+    links = {link.attrib["name"]: link for link in root.findall("link")}
+    children: dict[str, list[str]] = {}
+    parent_of: dict[str, str] = {}
+    edge = {}
+    for joint in root.findall("joint"):
+        parent = joint.find("parent").attrib["link"]
+        child = joint.find("child").attrib["link"]
+        origin = joint.find("origin")
+        axis_elem = joint.find("axis")
+        xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        axis = (
+            _parse_floats(axis_elem.attrib.get("xyz"), (1.0, 0.0, 0.0)) if axis_elem is not None else (1.0, 0.0, 0.0)
+        )
+        children.setdefault(parent, []).append(child)
+        parent_of[child] = parent
+        edge[(parent, child)] = (joint.attrib["name"], joint.attrib.get("type", "fixed"), xyz, rpy, axis)
+    return links, children, parent_of, edge
+
+
+def _link_world_transforms(root: ET.Element, joint_positions: dict[str, float]) -> tuple[dict[str, np.ndarray], dict[str, str], dict]:
+    links, children, parent_of, edge = _link_tree(root)
+    root_link = next(name for name in links if name not in parent_of)
+    world = {root_link: np.eye(4)}
+    stack = [root_link]
+    while stack:
+        parent = stack.pop()
+        for child in children.get(parent, []):
+            joint_name, joint_type, xyz, rpy, axis = edge[(parent, child)]
+            joint_transform = _transform(xyz, rpy)
+            if joint_type in {"revolute", "continuous"}:
+                rotation = np.eye(4)
+                rotation[:3, :3] = _axis_angle(axis, joint_positions.get(joint_name, 0.0))
+                joint_transform = joint_transform @ rotation
+            elif joint_type == "prismatic":
+                translation = np.eye(4)
+                translation[:3, 3] = np.asarray(axis, dtype=float) * joint_positions.get(joint_name, 0.0)
+                joint_transform = joint_transform @ translation
+            world[child] = world[parent] @ joint_transform
+            stack.append(child)
+    return world, parent_of, edge
+
+
+def _box_vertices(transform: np.ndarray, size: tuple[float, float, float]) -> np.ndarray:
+    sx, sy, sz = (np.asarray(size, dtype=float) * 0.5).tolist()
+    local = np.array(
+        [
+            [-sx, -sy, -sz, 1.0],
+            [sx, -sy, -sz, 1.0],
+            [sx, sy, -sz, 1.0],
+            [-sx, sy, -sz, 1.0],
+            [-sx, -sy, sz, 1.0],
+            [sx, -sy, sz, 1.0],
+            [sx, sy, sz, 1.0],
+            [-sx, sy, sz, 1.0],
+        ],
+        dtype=float,
+    )
+    return (transform @ local.T).T[:, :3]
+
+
+def _collect_boxes(root: ET.Element) -> tuple[list[Box], dict[str, str], dict]:
+    joint_positions = _load_default_joint_positions()
+    link_world, parent_of, edge = _link_world_transforms(root, joint_positions)
+    boxes = []
+    for link in root.findall("link"):
+        link_name = link.attrib["name"]
+        collision = link.find("collision")
+        box = link.find("./collision/geometry/box")
+        if collision is None or box is None:
+            continue
+        origin = collision.find("origin")
+        xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        size = _parse_floats(box.attrib["size"], (0.0, 0.0, 0.0))
+        transform = link_world[link_name] @ _transform(xyz, rpy)
+        boxes.append(Box(link_name, transform[:3, 3], transform[:3, :3], np.asarray(size, dtype=float) * 0.5, _box_vertices(transform, size)))
+    return boxes, parent_of, edge
+
+
+def _is_adjacent(parent_of: dict[str, str], a: str, b: str) -> bool:
+    return parent_of.get(a) == b or parent_of.get(b) == a
+
+
+def _obb_penetration(a: Box, b: Box) -> float | None:
+    axes = [a.axes[:, idx] for idx in range(3)] + [b.axes[:, idx] for idx in range(3)]
+    for i in range(3):
+        for j in range(3):
+            cross = np.cross(a.axes[:, i], b.axes[:, j])
+            norm = np.linalg.norm(cross)
+            if norm > 1.0e-8:
+                axes.append(cross / norm)
+
+    min_margin = float("inf")
+    for axis in axes:
+        axis = axis / np.linalg.norm(axis)
+        distance = abs(np.dot(b.center - a.center, axis))
+        radius_a = sum(a.half[i] * abs(np.dot(a.axes[:, i], axis)) for i in range(3))
+        radius_b = sum(b.half[i] * abs(np.dot(b.axes[:, i], axis)) for i in range(3))
+        margin = radius_a + radius_b - distance
+        if margin <= 0.0:
+            return None
+        min_margin = min(min_margin, margin)
+    return min_margin
+
+
+def _check_inertials(root: ET.Element):
+    total_mass = 0.0
+    missing_inertial = []
+    bad_mass = []
+    bad_inertia = []
+    for link in root.findall("link"):
+        name = link.attrib["name"]
+        inertial = link.find("inertial")
+        if inertial is None:
+            missing_inertial.append(name)
+            continue
+        mass = float(inertial.find("mass").attrib["value"])
+        total_mass += mass
+        if mass <= 0.0:
+            bad_mass.append((name, mass))
+        inertia = inertial.find("inertia").attrib
+        matrix = np.array(
+            [
+                [float(inertia.get("ixx", 0.0)), float(inertia.get("ixy", 0.0)), float(inertia.get("ixz", 0.0))],
+                [float(inertia.get("ixy", 0.0)), float(inertia.get("iyy", 0.0)), float(inertia.get("iyz", 0.0))],
+                [float(inertia.get("ixz", 0.0)), float(inertia.get("iyz", 0.0)), float(inertia.get("izz", 0.0))],
+            ]
+        )
+        eig = np.linalg.eigvalsh(matrix)
+        if eig.min() <= 0.0:
+            bad_inertia.append((name, eig.tolist()))
+    return total_mass, missing_inertial, bad_mass, bad_inertia
+
+
+def _check_joint_limits(root: ET.Element):
+    movable = []
+    missing = []
+    bad = []
+    for joint in root.findall("joint"):
+        if joint.attrib.get("type") == "fixed":
+            continue
+        name = joint.attrib["name"]
+        limit = joint.find("limit")
+        if limit is None:
+            missing.append(name)
+            continue
+        lower = float(limit.attrib.get("lower", "nan"))
+        upper = float(limit.attrib.get("upper", "nan"))
+        effort = float(limit.attrib.get("effort", "nan"))
+        velocity = float(limit.attrib.get("velocity", "nan"))
+        movable.append((name, lower, upper, effort, velocity))
+        if not (lower < upper and effort > 0.0 and velocity > 0.0):
+            bad.append((name, lower, upper, effort, velocity))
+    return movable, missing, bad
+
+
+def main() -> None:
+    root = ET.parse(URDF_PATH).getroot()
+    boxes, parent_of, edge = _collect_boxes(root)
+    movable, missing_limits, bad_limits = _check_joint_limits(root)
+    total_mass, missing_inertial, bad_mass, bad_inertia = _check_inertials(root)
+
+    adjacent_overlaps = []
+    non_adjacent_overlaps = []
+    for i, box_a in enumerate(boxes):
+        for box_b in boxes[i + 1 :]:
+            penetration = _obb_penetration(box_a, box_b)
+            if penetration is None:
+                continue
+            if _is_adjacent(parent_of, box_a.name, box_b.name):
+                joint_name = edge.get((box_a.name, box_b.name), edge.get((box_b.name, box_a.name), ("?",)))[0]
+                adjacent_overlaps.append((penetration, box_a.name, box_b.name, joint_name))
+            else:
+                non_adjacent_overlaps.append((penetration, box_a.name, box_b.name))
+
+    adjacent_overlaps.sort(reverse=True)
+    non_adjacent_overlaps.sort(reverse=True)
+
+    base_z = _load_base_init_z()
+    base_target = _load_base_height_target()
+    foot_min_z = {box.name: float(box.vertices[:, 2].min()) for box in boxes if box.name in FOOT_LINKS}
+    lowest_foot_z = min(foot_min_z.values())
+    suggested_base_z = -lowest_foot_z + 0.01
+
+    print("GEN2 RL READINESS CHECK")
+    print(f"urdf={URDF_PATH}")
+    print(f"links={len(root.findall('link'))} joints={len(root.findall('joint'))} collision_boxes={len(boxes)}")
+    print(f"total_mass_kg={total_mass:.6f}")
+    print(f"inertial_missing={len(missing_inertial)} bad_mass={len(bad_mass)} bad_inertia={len(bad_inertia)}")
+    print(f"movable_joints={len(movable)} missing_limits={len(missing_limits)} bad_limits={len(bad_limits)}")
+    if movable:
+        print(f"effort_range_nm={min(item[3] for item in movable):.3f}..{max(item[3] for item in movable):.3f}")
+        print(f"velocity_range_rad_s={min(item[4] for item in movable):.3f}..{max(item[4] for item in movable):.3f}")
+    print(f"non_adjacent_collision_overlaps={len(non_adjacent_overlaps)}")
+    for penetration, link_a, link_b in non_adjacent_overlaps:
+        print(f"  NON_ADJ {link_a} <-> {link_b}: penetration_m={penetration:.6f}")
+    print(f"adjacent_collision_overlaps={len(adjacent_overlaps)}")
+    for penetration, link_a, link_b, joint_name in adjacent_overlaps:
+        print(f"  ADJ {link_a} <-> {link_b} via {joint_name}: penetration_m={penetration:.6f}")
+    print("foot_collision_bottom_relative_to_body_root_m:")
+    for name, min_z in foot_min_z.items():
+        world_min = min_z + base_z if base_z is not None else float("nan")
+        print(f"  {name}: relative={min_z:.6f} world_at_init={world_min:.6f}")
+    print(f"base_init_z={base_z:.6f}" if base_z is not None else "base_init_z=unknown")
+    print(f"base_height_target={base_target:.6f}" if base_target is not None else "base_height_target=unknown")
+    print(f"suggested_base_z_for_1cm_foot_clearance={suggested_base_z:.6f}")
+
+    hard_fail = bool(missing_inertial or bad_mass or bad_inertia or missing_limits or bad_limits or non_adjacent_overlaps)
+    ground_penetration = base_z is not None and lowest_foot_z + base_z < -0.005
+    if hard_fail:
+        print("RESULT=FAIL")
+    elif ground_penetration:
+        print("RESULT=WARN_RESET_HEIGHT")
+    elif adjacent_overlaps:
+        print("RESULT=WARN_ADJACENT_OVERLAPS")
+    else:
+        print("RESULT=PASS")
+
+
+if __name__ == "__main__":
+    main()

scripts/gen2_generate_simplified_collisions.py

@@ -0,0 +1,292 @@
+#!/usr/bin/env python3
+"""Generate simplified Gen2 collision proxies and URDF variants.
+
+The script does not modify the original robot.urdf. It writes:
+
+- source/gen2_lab/assets/collision_simplified/*.stl
+- source/gen2_lab/assets/robot_simplified_collision.urdf
+- source/gen2_lab/assets/robot_simplified_collision_mesh.urdf
+- source/gen2_lab/assets/collision_simplified_audit.csv
+
+The recommended URDF is robot_simplified_collision.urdf. It uses URDF box
+primitives for collision. The generated STL files and mesh URDF are useful for
+visual overlay checks in tools such as Blender.
+"""
+
+from __future__ import annotations
+
+import csv
+import math
+import struct
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from xml.dom import minidom
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+ASSET_DIR = REPO_ROOT / "source" / "gen2_lab" / "assets"
+INPUT_URDF = ASSET_DIR / "robot.urdf"
+OUTPUT_DIR = ASSET_DIR / "collision_simplified"
+OUTPUT_PRIMITIVE_URDF = ASSET_DIR / "robot_simplified_collision.urdf"
+OUTPUT_MESH_URDF = ASSET_DIR / "robot_simplified_collision_mesh.urdf"
+AUDIT_CSV = ASSET_DIR / "collision_simplified_audit.csv"
+
+
+def _parse_floats(text: str | None, default: tuple[float, float, float]) -> tuple[float, float, float]:
+    if not text:
+        return default
+    values = [float(v) for v in text.split()]
+    if len(values) != 3:
+        raise ValueError(f"Expected 3 floats, got: {text}")
+    return values[0], values[1], values[2]
+
+
+def _rpy_to_matrix(rpy: tuple[float, float, float]) -> list[list[float]]:
+    roll, pitch, yaw = rpy
+    cr, sr = math.cos(roll), math.sin(roll)
+    cp, sp = math.cos(pitch), math.sin(pitch)
+    cy, sy = math.cos(yaw), math.sin(yaw)
+    return [
+        [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
+        [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
+        [-sp, cp * sr, cp * cr],
+    ]
+
+
+def _mat_vec_mul(mat: list[list[float]], vec: tuple[float, float, float]) -> tuple[float, float, float]:
+    return tuple(sum(mat[i][j] * vec[j] for j in range(3)) for i in range(3))
+
+
+def _fmt(values: tuple[float, ...] | list[float]) -> str:
+    return " ".join(f"{v:.9g}" for v in values)
+
+
+def _binary_stl_bbox(path: Path) -> tuple[tuple[float, float, float], tuple[float, float, float], int]:
+    data = path.read_bytes()
+    if len(data) < 84:
+        raise ValueError(f"STL too small: {path}")
+    tri_count = struct.unpack("<I", data[80:84])[0]
+    if 84 + 50 * tri_count != len(data):
+        raise ValueError(f"Only binary STL is supported by this script: {path}")
+
+    mins = [float("inf"), float("inf"), float("inf")]
+    maxs = [float("-inf"), float("-inf"), float("-inf")]
+    offset = 84
+    for _ in range(tri_count):
+        values = struct.unpack("<12f", data[offset : offset + 48])
+        for start in (3, 6, 9):
+            for axis in range(3):
+                value = values[start + axis]
+                mins[axis] = min(mins[axis], value)
+                maxs[axis] = max(maxs[axis], value)
+        offset += 50
+
+    extent = tuple(maxs[i] - mins[i] for i in range(3))
+    center = tuple((maxs[i] + mins[i]) * 0.5 for i in range(3))
+    return extent, center, tri_count
+
+
+def _link_scale(name: str) -> tuple[float, float, float]:
+    if name == "body_link":
+        return 0.74, 0.74, 0.74
+    if name.startswith("waist_link"):
+        return 0.22, 0.22, 0.22
+    if name.endswith("leg_link_6"):
+        return 0.90, 0.80, 0.45
+    if name.endswith("leg_link_4"):
+        return 0.40, 0.44, 0.58
+    if name.endswith("leg_link_5"):
+        return 0.42, 0.42, 0.50
+    if "leg_link_2" in name or "leg_link_3" in name:
+        return 0.30, 0.30, 0.30
+    if "leg_link_1" in name:
+        return 0.22, 0.22, 0.22
+    if "arm_link_7" in name:
+        return 0.20, 0.20, 0.20
+    if "arm_link_5" in name or "arm_link_6" in name:
+        return 0.22, 0.22, 0.22
+    if "arm_link_1" in name or "arm_link_2" in name:
+        return 0.30, 0.30, 0.30
+    if "arm_link" in name:
+        return 0.36, 0.36, 0.36
+    return 0.46, 0.46, 0.46
+
+
+def _proxy_box(name: str, extent: tuple[float, float, float], center: tuple[float, float, float]):
+    scale = _link_scale(name)
+    size = tuple(max(extent[i] * scale[i], 0.01) for i in range(3))
+    proxy_center = center
+
+    # Keep the sole close to the original mesh bottom while reducing foot height.
+    if name.endswith("leg_link_6"):
+        min_z = center[2] - extent[2] * 0.5
+        size = (max(extent[0] * 0.90, 0.08), max(extent[1] * 0.80, 0.04), min(max(extent[2] * 0.45, 0.025), 0.045))
+        proxy_center = (center[0], center[1], min_z + size[2] * 0.5 + 0.002)
+
+    return proxy_center, size, scale
+
+
+def _normal(a, b, c):
+    ux, uy, uz = b[0] - a[0], b[1] - a[1], b[2] - a[2]
+    vx, vy, vz = c[0] - a[0], c[1] - a[1], c[2] - a[2]
+    nx, ny, nz = uy * vz - uz * vy, uz * vx - ux * vz, ux * vy - uy * vx
+    length = math.sqrt(nx * nx + ny * ny + nz * nz)
+    if length <= 1e-12:
+        return 0.0, 0.0, 0.0
+    return nx / length, ny / length, nz / length
+
+
+def _write_binary_stl(path: Path, triangles: list[tuple[tuple[float, float, float], ...]]) -> None:
+    header = b"generated simplified collision proxy".ljust(80, b" ")
+    with path.open("wb") as file:
+        file.write(header)
+        file.write(struct.pack("<I", len(triangles)))
+        for tri in triangles:
+            n = _normal(*tri)
+            file.write(struct.pack("<3f", *n))
+            for vertex in tri:
+                file.write(struct.pack("<3f", *vertex))
+            file.write(struct.pack("<H", 0))
+
+
+def _box_triangles(center: tuple[float, float, float], size: tuple[float, float, float]):
+    cx, cy, cz = center
+    sx, sy, sz = (size[0] * 0.5, size[1] * 0.5, size[2] * 0.5)
+    v = [
+        (cx - sx, cy - sy, cz - sz),
+        (cx + sx, cy - sy, cz - sz),
+        (cx + sx, cy + sy, cz - sz),
+        (cx - sx, cy + sy, cz - sz),
+        (cx - sx, cy - sy, cz + sz),
+        (cx + sx, cy - sy, cz + sz),
+        (cx + sx, cy + sy, cz + sz),
+        (cx - sx, cy + sy, cz + sz),
+    ]
+    faces = [
+        (0, 2, 1),
+        (0, 3, 2),
+        (4, 5, 6),
+        (4, 6, 7),
+        (0, 1, 5),
+        (0, 5, 4),
+        (1, 2, 6),
+        (1, 6, 5),
+        (2, 3, 7),
+        (2, 7, 6),
+        (3, 0, 4),
+        (3, 4, 7),
+    ]
+    return [(v[a], v[b], v[c]) for a, b, c in faces]
+
+
+def _replace_collision_with_box(collision: ET.Element, origin_xyz, origin_rpy, size) -> None:
+    for child in list(collision):
+        collision.remove(child)
+
+    origin = ET.SubElement(collision, "origin")
+    origin.set("xyz", _fmt(origin_xyz))
+    origin.set("rpy", _fmt(origin_rpy))
+    geometry = ET.SubElement(collision, "geometry")
+    box = ET.SubElement(geometry, "box")
+    box.set("size", _fmt(size))
+
+
+def _replace_collision_with_mesh(collision: ET.Element, mesh_path: str) -> None:
+    geometry = collision.find("geometry")
+    if geometry is None:
+        geometry = ET.SubElement(collision, "geometry")
+    for child in list(geometry):
+        geometry.remove(child)
+    mesh = ET.SubElement(geometry, "mesh")
+    mesh.set("filename", mesh_path)
+
+
+def _write_pretty_xml(tree: ET.ElementTree, path: Path) -> None:
+    rough = ET.tostring(tree.getroot(), encoding="utf-8")
+    pretty = minidom.parseString(rough).toprettyxml(indent="  ")
+    lines = [line for line in pretty.splitlines() if line.strip()]
+    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
+
+
+def main() -> None:
+    if not INPUT_URDF.exists():
+        raise FileNotFoundError(INPUT_URDF)
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    for stale_proxy in OUTPUT_DIR.glob("*_box_collision.stl"):
+        stale_proxy.unlink()
+
+    primitive_tree = ET.parse(INPUT_URDF)
+    mesh_tree = ET.parse(INPUT_URDF)
+    audit_rows = []
+
+    mesh_links = {link.attrib["name"]: link for link in mesh_tree.getroot().findall("link")}
+
+    for primitive_link in primitive_tree.getroot().findall("link"):
+        name = primitive_link.attrib["name"]
+        collision = primitive_link.find("collision")
+        mesh_link = mesh_links[name]
+        mesh_collision = mesh_link.find("collision")
+        if collision is None or mesh_collision is None:
+            continue
+
+        mesh = collision.find("./geometry/mesh")
+        origin = collision.find("origin")
+        if mesh is None or origin is None:
+            continue
+
+        mesh_file = mesh.attrib["filename"]
+        extent, center, tri_count = _binary_stl_bbox(ASSET_DIR / mesh_file)
+        proxy_center, size, scale = _proxy_box(name, extent, center)
+
+        original_xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0))
+        original_rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0))
+        rotated_center = _mat_vec_mul(_rpy_to_matrix(original_rpy), proxy_center)
+        proxy_origin_xyz = tuple(original_xyz[i] + rotated_center[i] for i in range(3))
+
+        stl_name = f"{name}_box_collision.stl"
+        stl_rel_path = f"collision_simplified/{stl_name}"
+        _write_binary_stl(OUTPUT_DIR / stl_name, _box_triangles(proxy_center, size))
+
+        _replace_collision_with_box(collision, proxy_origin_xyz, original_rpy, size)
+        _replace_collision_with_mesh(mesh_collision, stl_rel_path)
+
+        audit_rows.append(
+            {
+                "link": name,
+                "source_mesh": mesh_file,
+                "proxy_mesh": stl_rel_path,
+                "source_triangles": tri_count,
+                "source_extent_x": extent[0],
+                "source_extent_y": extent[1],
+                "source_extent_z": extent[2],
+                "proxy_center_x": proxy_center[0],
+                "proxy_center_y": proxy_center[1],
+                "proxy_center_z": proxy_center[2],
+                "proxy_size_x": size[0],
+                "proxy_size_y": size[1],
+                "proxy_size_z": size[2],
+                "scale_x": scale[0],
+                "scale_y": scale[1],
+                "scale_z": scale[2],
+            }
+        )
+
+    if not audit_rows:
+        raise RuntimeError(f"No mesh collisions were converted from {INPUT_URDF}")
+
+    _write_pretty_xml(primitive_tree, OUTPUT_PRIMITIVE_URDF)
+    _write_pretty_xml(mesh_tree, OUTPUT_MESH_URDF)
+
+    with AUDIT_CSV.open("w", newline="", encoding="utf-8") as file:
+        writer = csv.DictWriter(file, fieldnames=list(audit_rows[0].keys()))
+        writer.writeheader()
+        writer.writerows(audit_rows)
+
+    print(f"wrote {OUTPUT_PRIMITIVE_URDF}")
+    print(f"wrote {OUTPUT_MESH_URDF}")
+    print(f"wrote {AUDIT_CSV}")
+    print(f"wrote {len(audit_rows)} collision STL files under {OUTPUT_DIR}")
+
+
+if __name__ == "__main__":
+    main()

scripts/gen2_visualize_collisions.py

@@ -0,0 +1,301 @@
+#!/usr/bin/env python3
+"""Render Gen2 collision boxes in the configured default standing pose."""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import math
+import re
+import xml.etree.ElementTree as ET
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_URDF = REPO_ROOT / "source" / "gen2_lab" / "assets" / "robot_simplified_collision.urdf"
+DEFAULT_OUTPUT = REPO_ROOT / "source" / "gen2_lab" / "assets" / "gen2_collision_visualization.png"
+GEN2_CFG = REPO_ROOT / "source" / "engineai_lab" / "robots" / "gen2.py"
+
+
+def _parse_floats(text: str | None, default: tuple[float, ...]) -> tuple[float, ...]:
+    if not text:
+        return default
+    return tuple(float(value) for value in text.split())
+
+
+def _rpy_matrix(rpy: tuple[float, float, float]) -> np.ndarray:
+    roll, pitch, yaw = rpy
+    cr, sr = math.cos(roll), math.sin(roll)
+    cp, sp = math.cos(pitch), math.sin(pitch)
+    cy, sy = math.cos(yaw), math.sin(yaw)
+    return np.array(
+        [
+            [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
+            [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
+            [-sp, cp * sr, cp * cr],
+        ],
+        dtype=float,
+    )
+
+
+def _axis_angle(axis: tuple[float, float, float], angle: float) -> np.ndarray:
+    vec = np.asarray(axis, dtype=float)
+    norm = np.linalg.norm(vec)
+    if norm < 1.0e-12:
+        return np.eye(3)
+    x, y, z = vec / norm
+    c, s = math.cos(angle), math.sin(angle)
+    one_c = 1.0 - c
+    return np.array(
+        [
+            [c + x * x * one_c, x * y * one_c - z * s, x * z * one_c + y * s],
+            [y * x * one_c + z * s, c + y * y * one_c, y * z * one_c - x * s],
+            [z * x * one_c - y * s, z * y * one_c + x * s, c + z * z * one_c],
+        ],
+        dtype=float,
+    )
+
+
+def _transform(xyz=(0.0, 0.0, 0.0), rpy=(0.0, 0.0, 0.0)) -> np.ndarray:
+    matrix = np.eye(4)
+    matrix[:3, :3] = _rpy_matrix(rpy)
+    matrix[:3, 3] = np.asarray(xyz, dtype=float)
+    return matrix
+
+
+def _load_default_joint_positions() -> dict[str, float]:
+    if not GEN2_CFG.exists():
+        return {}
+    text = GEN2_CFG.read_text(encoding="utf-8")
+    match = re.search(r"joint_pos=\{(.*?)\},\n\s*joint_vel=", text, flags=re.S)
+    if match is None:
+        return {}
+    return ast.literal_eval("{" + match.group(1) + "}")
+
+
+def _compute_link_world_transforms(root: ET.Element, joint_positions: dict[str, float]) -> dict[str, np.ndarray]:
+    link_names = [link.attrib["name"] for link in root.findall("link")]
+    children: dict[str, list[str]] = {}
+    parent_of: dict[str, str] = {}
+    edge = {}
+
+    for joint in root.findall("joint"):
+        parent = joint.find("parent").attrib["link"]
+        child = joint.find("child").attrib["link"]
+        origin = joint.find("origin")
+        axis_elem = joint.find("axis")
+        xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        axis = (
+            _parse_floats(axis_elem.attrib.get("xyz"), (1.0, 0.0, 0.0)) if axis_elem is not None else (1.0, 0.0, 0.0)
+        )
+        children.setdefault(parent, []).append(child)
+        parent_of[child] = parent
+        edge[(parent, child)] = (joint.attrib["name"], joint.attrib.get("type", "fixed"), xyz, rpy, axis)
+
+    root_link = next(name for name in link_names if name not in parent_of)
+    world = {root_link: np.eye(4)}
+    stack = [root_link]
+    while stack:
+        parent = stack.pop()
+        for child in children.get(parent, []):
+            joint_name, joint_type, xyz, rpy, axis = edge[(parent, child)]
+            joint_transform = _transform(xyz, rpy)
+            if joint_type in {"revolute", "continuous"}:
+                rotation = np.eye(4)
+                rotation[:3, :3] = _axis_angle(axis, joint_positions.get(joint_name, 0.0))
+                joint_transform = joint_transform @ rotation
+            elif joint_type == "prismatic":
+                translation = np.eye(4)
+                translation[:3, 3] = np.asarray(axis, dtype=float) * joint_positions.get(joint_name, 0.0)
+                joint_transform = joint_transform @ translation
+            world[child] = world[parent] @ joint_transform
+            stack.append(child)
+    return world
+
+
+def _box_vertices(transform: np.ndarray, size: tuple[float, float, float]) -> np.ndarray:
+    sx, sy, sz = (np.asarray(size, dtype=float) * 0.5).tolist()
+    local = np.array(
+        [
+            [-sx, -sy, -sz, 1.0],
+            [sx, -sy, -sz, 1.0],
+            [sx, sy, -sz, 1.0],
+            [-sx, sy, -sz, 1.0],
+            [-sx, -sy, sz, 1.0],
+            [sx, -sy, sz, 1.0],
+            [sx, sy, sz, 1.0],
+            [-sx, sy, sz, 1.0],
+        ]
+    )
+    return (transform @ local.T).T[:, :3]
+
+
+def _box_faces(vertices: np.ndarray) -> list[np.ndarray]:
+    indices = [(0, 1, 2, 3), (4, 5, 6, 7), (0, 1, 5, 4), (1, 2, 6, 5), (2, 3, 7, 6), (3, 0, 4, 7)]
+    return [vertices[list(face)] for face in indices]
+
+
+def _module_color(link_name: str) -> str:
+    if link_name.startswith("left_leg"):
+        return "#2a9d8f"
+    if link_name.startswith("right_leg"):
+        return "#e76f51"
+    if link_name.startswith("left_arm"):
+        return "#457b9d"
+    if link_name.startswith("right_arm"):
+        return "#f4a261"
+    if link_name.startswith("waist"):
+        return "#8d5a97"
+    return "#5c677d"
+
+
+def _collect_boxes(root: ET.Element) -> list[dict]:
+    joint_positions = _load_default_joint_positions()
+    link_world = _compute_link_world_transforms(root, joint_positions)
+    boxes = []
+    for link in root.findall("link"):
+        link_name = link.attrib["name"]
+        collision = link.find("collision")
+        box = link.find("./collision/geometry/box")
+        if collision is None or box is None:
+            continue
+        origin = collision.find("origin")
+        xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
+        size = _parse_floats(box.attrib["size"], (0.0, 0.0, 0.0))
+        transform = link_world[link_name] @ _transform(xyz, rpy)
+        boxes.append(
+            {
+                "name": link_name,
+                "vertices": _box_vertices(transform, size),
+                "center": transform[:3, 3],
+                "size": size,
+                "color": _module_color(link_name),
+            }
+        )
+    return boxes
+
+
+def _set_equal_axes_3d(ax, points: np.ndarray) -> None:
+    mins = points.min(axis=0)
+    maxs = points.max(axis=0)
+    center = (mins + maxs) * 0.5
+    radius = float((maxs - mins).max() * 0.56)
+    ax.set_xlim(center[0] - radius, center[0] + radius)
+    ax.set_ylim(center[1] - radius, center[1] + radius)
+    ax.set_zlim(center[2] - radius, center[2] + radius)
+
+
+def _draw_projection(ax, boxes: list[dict], dims: tuple[int, int], labels: tuple[str, str], title: str) -> None:
+    edges = [(0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (5, 6), (6, 7), (7, 4), (0, 4), (1, 5), (2, 6), (3, 7)]
+    for box in boxes:
+        vertices = box["vertices"]
+        for start, end in edges:
+            ax.plot(
+                [vertices[start, dims[0]], vertices[end, dims[0]]],
+                [vertices[start, dims[1]], vertices[end, dims[1]]],
+                color=box["color"],
+                linewidth=1.2,
+            )
+    ax.set_title(title)
+    ax.set_xlabel(labels[0])
+    ax.set_ylabel(labels[1])
+    ax.set_aspect("equal", adjustable="box")
+    ax.grid(True, alpha=0.25)
+
+
+def _save_projection(output_path: Path, boxes: list[dict], dims: tuple[int, int], labels: tuple[str, str], title: str) -> None:
+    fig, ax = plt.subplots(figsize=(8, 8), dpi=180)
+    _draw_projection(ax, boxes, dims, labels, title)
+    fig.tight_layout()
+    fig.savefig(output_path)
+    plt.close(fig)
+    print(f"wrote {output_path}")
+
+
+def _save_3d(output_path: Path, boxes: list[dict], all_points: np.ndarray) -> None:
+    fig = plt.figure(figsize=(8, 8), dpi=180)
+    ax3d = fig.add_subplot(1, 1, 1, projection="3d")
+    for box in boxes:
+        collection = Poly3DCollection(_box_faces(box["vertices"]), facecolor=box["color"], edgecolor="#202020", alpha=0.32)
+        collection.set_linewidth(0.45)
+        ax3d.add_collection3d(collection)
+    _set_equal_axes_3d(ax3d, all_points)
+    ax3d.set_title("Gen2 collision boxes - 3D")
+    ax3d.set_xlabel("X forward (m)")
+    ax3d.set_ylabel("Y left (m)")
+    ax3d.set_zlabel("Z up (m)")
+    ax3d.view_init(elev=18, azim=-58)
+    fig.tight_layout()
+    fig.savefig(output_path)
+    plt.close(fig)
+    print(f"wrote {output_path}")
+
+
+def render(urdf_path: Path, output_path: Path) -> None:
+    root = ET.parse(urdf_path).getroot()
+    boxes = _collect_boxes(root)
+    if not boxes:
+        raise RuntimeError(f"No box collisions found in {urdf_path}")
+
+    all_points = np.concatenate([box["vertices"] for box in boxes], axis=0)
+    fig = plt.figure(figsize=(15, 11), dpi=170)
+
+    ax3d = fig.add_subplot(2, 2, 1, projection="3d")
+    for box in boxes:
+        collection = Poly3DCollection(_box_faces(box["vertices"]), facecolor=box["color"], edgecolor="#202020", alpha=0.32)
+        collection.set_linewidth(0.45)
+        ax3d.add_collection3d(collection)
+    _set_equal_axes_3d(ax3d, all_points)
+    ax3d.set_title("Gen2 collision boxes - 3D")
+    ax3d.set_xlabel("X forward (m)")
+    ax3d.set_ylabel("Y left (m)")
+    ax3d.set_zlabel("Z up (m)")
+    ax3d.view_init(elev=18, azim=-58)
+
+    _draw_projection(fig.add_subplot(2, 2, 2), boxes, (1, 2), ("Y left (m)", "Z up (m)"), "Front view")
+    _draw_projection(fig.add_subplot(2, 2, 3), boxes, (0, 2), ("X forward (m)", "Z up (m)"), "Side view")
+    _draw_projection(fig.add_subplot(2, 2, 4), boxes, (0, 1), ("X forward (m)", "Y left (m)"), "Top view")
+
+    legend_items = [
+        ("body", "#5c677d"),
+        ("waist", "#8d5a97"),
+        ("left leg", "#2a9d8f"),
+        ("right leg", "#e76f51"),
+        ("left arm", "#457b9d"),
+        ("right arm", "#f4a261"),
+    ]
+    handles = [plt.Line2D([0], [0], color=color, lw=4, label=name) for name, color in legend_items]
+    fig.legend(handles=handles, loc="lower center", ncol=6, frameon=False)
+    fig.suptitle(f"{urdf_path.name}: {len(boxes)} collision boxes", fontsize=14)
+    fig.tight_layout(rect=(0, 0.04, 1, 0.96))
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output_path)
+    plt.close(fig)
+    print(f"wrote {output_path}")
+
+    stem = output_path.with_suffix("")
+    _save_3d(stem.with_name(f"{stem.name}_3d.png"), boxes, all_points)
+    _save_projection(stem.with_name(f"{stem.name}_front.png"), boxes, (1, 2), ("Y left (m)", "Z up (m)"), "Front view")
+    _save_projection(stem.with_name(f"{stem.name}_side.png"), boxes, (0, 2), ("X forward (m)", "Z up (m)"), "Side view")
+    _save_projection(stem.with_name(f"{stem.name}_top.png"), boxes, (0, 1), ("X forward (m)", "Y left (m)"), "Top view")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--urdf", type=Path, default=DEFAULT_URDF)
+    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
+    args = parser.parse_args()
+    render(args.urdf.resolve(), args.output.resolve())
+
+
+if __name__ == "__main__":
+    main()

scripts/play.py

@@ -0,0 +1,227 @@
+"""Script to play a checkpoint from an RSL-RL agent."""
+
+"""Launch Isaac Sim Simulator first."""
+
+import argparse
+import sys
+from pathlib import Path
+
+from isaaclab.app import AppLauncher
+
+# local imports
+import cli_args  # isort: skip
+
+# ensure repository root is on Python path for Hydra registry imports
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.append(str(REPO_ROOT))
+
+# add argparse arguments
+parser = argparse.ArgumentParser(description="Play an RSL-RL policy checkpoint.")
+parser.add_argument(
+    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
+)
+parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
+parser.add_argument("--task", type=str, default=None, help="Name of the task.")
+parser.add_argument("--max_steps", type=int, default=None, help="Stop replay after this many policy steps.")
+
+command_group = parser.add_argument_group("velocity command", description="Velocity command source during replay.")
+command_group.add_argument(
+    "--command_source",
+    type=str,
+    choices=("random", "fixed", "keyboard"),
+    default="random",
+    help="Use environment-generated, fixed, or keyboard velocity commands.",
+)
+command_group.add_argument("--vx", type=float, default=0.3, help="Fixed forward velocity in m/s.")
+command_group.add_argument("--vy", type=float, default=0.0, help="Fixed lateral velocity in m/s.")
+command_group.add_argument("--wz", type=float, default=0.0, help="Fixed yaw velocity in rad/s.")
+command_group.add_argument(
+    "--command_mode",
+    type=str,
+    choices=("step", "ramp"),
+    default="ramp",
+    help="Apply external commands immediately or through acceleration limits.",
+)
+command_group.add_argument("--linear_accel", type=float, default=0.8, help="Planar acceleration limit in m/s^2.")
+command_group.add_argument("--yaw_accel", type=float, default=1.5, help="Yaw acceleration limit in rad/s^2.")
+command_group.add_argument("--keyboard_vx", type=float, default=0.4, help="Keyboard forward velocity in m/s.")
+command_group.add_argument("--keyboard_vy", type=float, default=0.2, help="Keyboard lateral velocity in m/s.")
+command_group.add_argument("--keyboard_wz", type=float, default=0.5, help="Keyboard yaw velocity in rad/s.")
+
+# append RSL-RL cli arguments
+cli_args.add_rsl_rl_args(parser)
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+args_cli, hydra_args = parser.parse_known_args()
+if args_cli.linear_accel <= 0.0:
+    parser.error("--linear_accel must be positive.")
+if args_cli.yaw_accel <= 0.0:
+    parser.error("--yaw_accel must be positive.")
+if args_cli.max_steps is not None and args_cli.max_steps <= 0:
+    parser.error("--max_steps must be positive.")
+# always enable cameras to record video
+# clear out sys.argv for Hydra
+sys.argv = [sys.argv[0]] + hydra_args
+
+# launch omniverse app
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import os
+
+import gymnasium as gym
+import torch
+
+from rsl_rl.runners import OnPolicyRunner
+
+from isaaclab.envs import (
+    DirectMARLEnv,
+    DirectMARLEnvCfg,
+    DirectRLEnvCfg,
+    ManagerBasedRLEnvCfg,
+    multi_agent_to_single_agent,
+)
+from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
+from isaaclab_tasks.utils import get_checkpoint_path
+from isaaclab_tasks.utils.hydra import hydra_task_config
+
+# Import extensions to set up environment tasks
+import engineai_lab.tasks  # noqa: F401
+
+
+class VelocityCommandController:
+    """Apply fixed or keyboard commands directly to a velocity command term."""
+
+    def __init__(self, env: RslRlVecEnvWrapper, args: argparse.Namespace):
+        self.env = env
+        self.args = args
+        self.command_term = env.unwrapped.command_manager.get_term("base_velocity")
+        if not hasattr(self.command_term, "vel_command_b"):
+            raise TypeError("The base_velocity command term does not expose a velocity command buffer.")
+
+        self.current_command = torch.zeros(3, device=env.unwrapped.device)
+        self.fixed_command = torch.tensor([args.vx, args.vy, args.wz], device=env.unwrapped.device)
+        self.keyboard = None
+
+        if args.command_source == "keyboard":
+            from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
+
+            keyboard_cfg = Se2KeyboardCfg(
+                sim_device=str(env.unwrapped.device),
+                v_x_sensitivity=args.keyboard_vx,
+                v_y_sensitivity=args.keyboard_vy,
+                omega_z_sensitivity=args.keyboard_wz,
+            )
+            self.keyboard = Se2Keyboard(keyboard_cfg)
+            print(self.keyboard)
+        else:
+            print(
+                "[INFO] Fixed velocity command: "
+                f"vx={args.vx:.3f} m/s, vy={args.vy:.3f} m/s, wz={args.wz:.3f} rad/s"
+            )
+        print(f"[INFO] External command mode: {args.command_mode}")
+
+    def update(self):
+        target = self.keyboard.advance() if self.keyboard is not None else self.fixed_command
+        if self.args.command_mode == "step":
+            self.current_command.copy_(target)
+        else:
+            self._apply_ramp(target)
+
+        self.command_term.vel_command_b[:] = self.current_command.unsqueeze(0)
+        self.command_term.time_left.fill_(float("inf"))
+        self.command_term.is_standing_env.fill_(False)
+        if hasattr(self.command_term, "is_heading_env"):
+            self.command_term.is_heading_env.fill_(False)
+
+    def _apply_ramp(self, target: torch.Tensor):
+        dt = self.env.unwrapped.step_dt
+        linear_delta = target[:2] - self.current_command[:2]
+        linear_delta_norm = torch.linalg.norm(linear_delta)
+        max_linear_delta = self.args.linear_accel * dt
+        linear_scale = torch.clamp(max_linear_delta / torch.clamp(linear_delta_norm, min=1.0e-6), max=1.0)
+        linear_delta = linear_delta * linear_scale
+        self.current_command[:2] += linear_delta
+
+        yaw_delta = torch.clamp(
+            target[2] - self.current_command[2],
+            min=-self.args.yaw_accel * dt,
+            max=self.args.yaw_accel * dt,
+        )
+        self.current_command[2] += yaw_delta
+
+
+@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point")
+def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
+    """Play with RSL-RL agent."""
+    agent_cfg: RslRlOnPolicyRunnerCfg = cli_args.parse_rsl_rl_cfg(args_cli.task, args_cli)
+    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
+
+    # specify directory for logging experiments
+    log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
+    log_root_path = os.path.abspath(log_root_path)
+    print(f"[INFO] Loading experiment from directory: {log_root_path}")
+    resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
+
+    log_dir = os.path.dirname(resume_path)
+
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
+    # create isaac environment
+    env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
+
+    # convert to single-agent instance if required by the RL algorithm
+    if isinstance(env.unwrapped, DirectMARLEnv):
+        env = multi_agent_to_single_agent(env)
+
+    # wrap around environment for rsl-rl
+    env = RslRlVecEnvWrapper(env)
+
+    command_controller = None
+    if args_cli.command_source != "random":
+        command_controller = VelocityCommandController(env, args_cli)
+
+    # load previously trained model
+    ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
+    ppo_runner.load(resume_path)
+
+    # obtain the trained policy for inference
+    policy = ppo_runner.get_inference_policy(device=env.unwrapped.device)
+
+    # export policy to onnx/jit
+    export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
+    file_basename = os.path.basename(resume_path).split(".")[0]
+    ppo_runner.export_policy_to_onnx(path=export_model_dir, filename=file_basename+".onnx")
+
+    # reset environment
+    obs = env.get_observations()
+    step_count = 0
+    # simulate environment
+    while simulation_app.is_running() and (args_cli.max_steps is None or step_count < args_cli.max_steps):
+        # run everything in inference mode
+        with torch.inference_mode():
+            if command_controller is not None:
+                command_controller.update()
+                obs = env.get_observations()
+            # agent stepping
+            actions = policy(obs)
+            # env stepping
+            obs, _, _, _ = env.step(actions)
+            step_count += 1
+
+    if args_cli.max_steps is not None:
+        print(f"[INFO] Replay completed after {step_count} policy steps.")
+
+    # close the simulator
+    env.close()
+
+
+if __name__ == "__main__":
+    # run the main function
+    main()
+    # close sim app
+    simulation_app.close()

scripts/replay_npz.py

@@ -0,0 +1,169 @@
+"""This script demonstrates how to use the interactive scene interface to setup a scene with multiple prims.
+
+.. code-block:: bash
+
+    # Usage
+    python replay_motion.py --motion_file dataset/xxxx.npz
+"""
+
+"""Launch Isaac Sim Simulator first."""
+
+import argparse
+import os
+import numpy as np
+import torch
+from typing import Sequence
+
+from isaaclab.app import AppLauncher
+
+# add argparse arguments
+parser = argparse.ArgumentParser(description="Replay converted motions.")
+parser.add_argument("--motion_file", type=str, required=True, help="Path to the motion file (.npz).")
+
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+# parse the arguments
+args_cli = parser.parse_args()
+
+# launch omniverse app
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import isaaclab.sim as sim_utils
+from isaaclab.assets import Articulation, ArticulationCfg, AssetBaseCfg
+from isaaclab.scene import InteractiveScene, InteractiveSceneCfg
+from isaaclab.sim import SimulationContext
+from isaaclab.utils import configclass
+from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
+
+# ! terrain
+from isaaclab.terrains import TerrainImporterCfg
+
+
+##
+# Pre-defined configs
+##
+from engineai_lab.robots.pm01 import PM01_CFG
+
+class MotionLoader:
+    def __init__(self, motion_file: str, body_indexes: Sequence[int], device: str = "cpu"):
+        assert os.path.isfile(motion_file), f"Invalid file path: {motion_file}"
+        data = np.load(motion_file)
+
+        self.fps = data["fps"]
+        self.joint_pos = torch.tensor(data["joint_pos"], dtype=torch.float32, device=device)
+        self.joint_vel = torch.tensor(data["joint_vel"], dtype=torch.float32, device=device)
+        self._body_pos_w = torch.tensor(data["body_pos_w"], dtype=torch.float32, device=device)
+        self._body_quat_w = torch.tensor(data["body_quat_w"], dtype=torch.float32, device=device)
+        self._body_lin_vel_w = torch.tensor(data["body_lin_vel_w"], dtype=torch.float32, device=device)
+        self._body_ang_vel_w = torch.tensor(data["body_ang_vel_w"], dtype=torch.float32, device=device)
+        self._body_indexes = body_indexes
+        self.time_step_total = self.joint_pos.shape[0]
+
+    @property
+    def body_pos_w(self) -> torch.Tensor:
+        return self._body_pos_w[:, self._body_indexes]
+
+    @property
+    def body_quat_w(self) -> torch.Tensor:
+        return self._body_quat_w[:, self._body_indexes]
+
+    @property
+    def body_lin_vel_w(self) -> torch.Tensor:
+        return self._body_lin_vel_w[:, self._body_indexes]
+
+    @property
+    def body_ang_vel_w(self) -> torch.Tensor:
+        return self._body_ang_vel_w[:, self._body_indexes]
+
+
+@configclass
+class ReplayMotionsSceneCfg(InteractiveSceneCfg):
+    """Configuration for a replay motions scene."""
+
+    terrain = TerrainImporterCfg(
+        prim_path="/World/ground",
+        terrain_type="plane",
+        collision_group=-1,
+        physics_material=sim_utils.RigidBodyMaterialCfg(
+            friction_combine_mode="multiply",
+            restitution_combine_mode="multiply",
+            static_friction=1.0,
+            dynamic_friction=1.0,
+        ),
+        visual_material=sim_utils.MdlFileCfg(
+            mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl",
+            project_uvw=True,
+        ),
+    )
+
+    sky_light = AssetBaseCfg(
+        prim_path="/World/skyLight",
+        spawn=sim_utils.DomeLightCfg(
+            intensity=750.0,
+            texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
+        ),
+    )
+
+    # articulation
+    robot: ArticulationCfg = PM01_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+
+
+def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene, motion_loader: MotionLoader):
+    # Extract scene entities
+    robot: Articulation = scene["robot"]
+    # Define simulation stepping
+    sim_dt = sim.get_physics_dt()
+
+    motion = motion_loader
+    time_steps = torch.zeros(scene.num_envs, dtype=torch.long, device=sim.device)
+
+    # Simulation loop
+    while simulation_app.is_running():
+        time_steps += 1
+        reset_ids = time_steps >= motion.time_step_total
+        time_steps[reset_ids] = 0
+
+        root_states = robot.data.default_root_state.clone()
+        root_states[:, :3] = motion.body_pos_w[time_steps][:, 0] + scene.env_origins[:, None, :]
+        root_states[:, 3:7] = motion.body_quat_w[time_steps][:, 0]
+        root_states[:, 7:10] = motion.body_lin_vel_w[time_steps][:, 0]
+        root_states[:, 10:] = motion.body_ang_vel_w[time_steps][:, 0]
+
+        robot.write_root_state_to_sim(root_states)
+        robot.write_joint_state_to_sim(motion.joint_pos[time_steps], motion.joint_vel[time_steps])
+        scene.write_data_to_sim()
+        sim.render()  # We don't want physic (sim.step())
+        scene.update(sim_dt)
+
+        pos_lookat = root_states[0, :3].cpu().numpy()
+        sim.set_camera_view(pos_lookat + np.array([2.0, 2.0, 0.5]), pos_lookat)
+
+
+def main():
+    sim_cfg = sim_utils.SimulationCfg(device=args_cli.device)
+    sim_cfg.dt = 0.01
+    sim = SimulationContext(sim_cfg)
+
+    motion_loader = MotionLoader(
+        motion_file=args_cli.motion_file,
+        body_indexes=torch.tensor([0], dtype=torch.long, device=sim.device),
+        device=sim.device,
+    )
+
+    scene_cfg = ReplayMotionsSceneCfg(num_envs=1, env_spacing=2.0)
+
+    scene = InteractiveScene(scene_cfg)
+
+    sim.reset()
+    # Run the simulator
+    run_simulator(sim, scene, motion_loader)
+
+
+if __name__ == "__main__":
+    # run the main function
+    main()
+    # close sim app
+    simulation_app.close()

scripts/train.py

@@ -0,0 +1,156 @@
+"""Script to train RL agent with RSL-RL."""
+
+"""Launch Isaac Sim Simulator first."""
+
+import argparse
+import sys
+from pathlib import Path
+
+from isaaclab.app import AppLauncher
+
+# local imports
+import cli_args  # isort: skip
+
+# ensure repository root is on the Python path for Hydra registry imports
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.append(str(REPO_ROOT))
+
+# add argparse arguments
+parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
+parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
+parser.add_argument("--task", type=str, default=None, help="Name of the task.")
+parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
+parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
+parser.add_argument(
+    "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes."
+)
+
+# append RSL-RL cli arguments
+cli_args.add_rsl_rl_args(parser)
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+args_cli, hydra_args = parser.parse_known_args()
+
+
+# clear out sys.argv for Hydra
+sys.argv = [sys.argv[0]] + hydra_args
+
+# launch omniverse app
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import gymnasium as gym
+import os
+import torch
+from datetime import datetime
+
+from isaaclab.envs import (
+    DirectMARLEnv,
+    DirectMARLEnvCfg,
+    DirectRLEnvCfg,
+    ManagerBasedRLEnvCfg,
+    multi_agent_to_single_agent,
+)
+from isaaclab.utils.dict import print_dict
+from isaaclab.utils.io import dump_yaml
+from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
+from isaaclab_tasks.utils import get_checkpoint_path
+from isaaclab_tasks.utils.hydra import hydra_task_config
+
+# Import extensions to set up environment tasks
+import engineai_lab.tasks  # noqa: F401
+from rsl_rl.runners.on_policy_runner import OnPolicyRunner
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.deterministic = False
+torch.backends.cudnn.benchmark = False
+
+
+@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point")
+def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
+    """Train with RSL-RL agent."""
+    # override configurations with non-hydra CLI arguments
+    agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
+    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
+    agent_cfg.max_iterations = (
+        args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations
+    )
+
+    # set the environment seed
+    # note: certain randomizations occur in the environment initialization so we set the seed here
+    env_cfg.seed = agent_cfg.seed
+    env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
+    if args_cli.distributed and args_cli.device is not None and "cpu" in args_cli.device:
+        raise ValueError(
+            "Distributed training is not supported when using CPU device. "
+            "Please use GPU device (e.g. --device cuda) for distributed training."
+        )
+
+    if args_cli.distributed:
+        env_cfg.sim.device = f"cuda:{app_launcher.local_rank}"
+        agent_cfg.device = f"cuda:{app_launcher.local_rank}"
+
+        seed = agent_cfg.seed + app_launcher.local_rank
+        env_cfg.seed = seed
+        agent_cfg.seed = seed
+
+    # specify directory for logging experiments
+    log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
+    log_root_path = os.path.abspath(log_root_path)
+    print(f"[INFO] Logging experiment in directory: {log_root_path}")
+    # specify directory for logging runs: {time-stamp}_{run_name}
+    log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    if agent_cfg.run_name:
+        log_dir += f"_{agent_cfg.run_name}"
+    log_dir = os.path.join(log_root_path, log_dir)
+
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
+    # create isaac environment
+    env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
+    # wrap for video recording
+
+    # convert to single-agent instance if required by the RL algorithm
+    if isinstance(env.unwrapped, DirectMARLEnv):
+        env = multi_agent_to_single_agent(env)
+
+    # wrap around environment for rsl-rl
+    env = RslRlVecEnvWrapper(env)
+
+    # create runner from rsl-rl
+    runner = OnPolicyRunner(
+        env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device
+    )
+    # write git state to logs
+    runner.add_git_repo_to_log(__file__)
+    # save resume path before creating a new log_dir
+    if agent_cfg.resume:
+        # get path to previous checkpoint
+        resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
+        print(f"[INFO]: Loading model checkpoint from: {resume_path}")
+        # load previously trained model
+        runner.load(resume_path)
+
+    # dump the configuration into log-directory
+    dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
+    dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)
+    # dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg)
+    # dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg)
+
+    # run training
+    runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)
+
+    # close the simulator
+    env.close()
+
+
+if __name__ == "__main__":
+    # run the main function
+    main()
+    # close sim app
+    simulation_app.close()

setup.py

@@ -0,0 +1,42 @@
+"""Installation script for the 'engineai_lab' python package."""
+
+import os
+import toml
+
+from setuptools import setup, find_packages
+
+# Obtain the extension data from the extension.toml file
+EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__))
+# Read the extension.toml file
+EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml"))
+
+# Minimum dependencies required prior to installation
+INSTALL_REQUIRES = [
+    "psutil",
+    "onnxscript",
+    "wandb>=0.19",
+]
+
+# Installation operation
+setup(
+    name="engineai_lab",
+    package_dir={"": "source"},
+    packages=find_packages(where="source"),
+    author=EXTENSION_TOML_DATA["package"]["author"],
+    maintainer=EXTENSION_TOML_DATA["package"]["maintainer"],
+    url=EXTENSION_TOML_DATA["package"]["repository"],
+    version=EXTENSION_TOML_DATA["package"]["version"],
+    description=EXTENSION_TOML_DATA["package"]["description"],
+    keywords=EXTENSION_TOML_DATA["package"]["keywords"],
+    install_requires=INSTALL_REQUIRES,
+    license="MIT",
+    include_package_data=True,
+    python_requires=">=3.10",
+    classifiers=[
+        "Natural Language :: English",
+        "Programming Language :: Python :: 3.10",
+        "Isaac Sim :: 2023.1.1",
+        "Isaac Sim :: 4.0.0",
+    ],
+    zip_safe=False,
+)

source/engineai_lab/__init__.py

@@ -0,0 +1,9 @@
+"""
+Python module serving as a project/extension template.
+"""
+
+# Register Gym environments.
+from .tasks import *
+
+# Ensure AMP data loaders register themselves on import.
+from .utils import AMP_data_loader  # noqa: F401

source/engineai_lab/algorithms/amp_ppo.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from tensordict import TensorDict
+
+from rsl_rl.env import VecEnv
+from rsl_rl.models import MLPModel
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.algorithms import PPO
+
+from engineai_lab.utils.AMP_discriminator import Discriminator
+from engineai_lab.utils.AMP_data_loader import AMPDataLoader
+
+
+
+class AMPPPO(PPO):
+    discriminator: Discriminator
+
+    def __init__(
+        self,
+        actor: MLPModel,
+        critic: MLPModel,
+        storage: RolloutStorage,
+        num_learning_epochs: int = 5,
+        num_mini_batches: int = 4,
+        clip_param: float = 0.2,
+        gamma: float = 0.99,
+        lam: float = 0.95,
+        value_loss_coef: float = 1.0,
+        entropy_coef: float = 0.01,
+        learning_rate: float = 0.001,
+        max_grad_norm: float = 1.0,
+        optimizer: str = "adam",
+        use_clipped_value_loss: bool = True,
+        schedule: str = "adaptive",
+        desired_kl: float = 0.01,
+        normalize_advantage_per_mini_batch: bool = False,
+        device: str = "cpu",
+        #  AMP parameters
+        discriminator: MLPModel = None,
+        data_loader: AMPDataLoader = None,
+        style_reward_weight: float = 2.0,
+        # RND parameters
+        rnd_cfg: dict | None = None,
+        # Symmetry parameters
+        symmetry_cfg: dict | None = None,
+        # Distributed training parameters
+        multi_gpu_cfg: dict | None = None,
+    ) -> None:
+
+        self.style_reward_weight = style_reward_weight
+        print(f"Initialized AMPPPO with style reward weight: {self.style_reward_weight}")
+
+        self.discriminator = discriminator
+        if self.discriminator is None:
+            raise ValueError("Discriminator must be provided for AMPPPO.")
+    
+        self.discriminator_data_loader = data_loader
+        if self.discriminator_data_loader is None:
+            raise ValueError("Data loader must be provided for AMPPPO.")
+    
+        super().__init__(
+            actor=actor,
+            critic=critic,
+            storage=storage,
+            num_learning_epochs=num_learning_epochs,
+            num_mini_batches=num_mini_batches,
+            clip_param=clip_param,
+            gamma=gamma,
+            lam=lam,
+            value_loss_coef=value_loss_coef,
+            entropy_coef=entropy_coef,
+            learning_rate=learning_rate,
+            max_grad_norm=max_grad_norm,
+            optimizer=optimizer,
+            use_clipped_value_loss=use_clipped_value_loss,
+            schedule=schedule,
+            desired_kl=desired_kl,
+            normalize_advantage_per_mini_batch=normalize_advantage_per_mini_batch,
+            device=device,
+            # RND parameters
+            rnd_cfg=rnd_cfg,
+            # Symmetry parameters
+            symmetry_cfg=symmetry_cfg,
+            # Distributed training parameters
+            multi_gpu_cfg=multi_gpu_cfg,
+        )
+
+        self.disc_optimizer = optim.Adam(self.discriminator.parameters(), lr=1e-4)
+
+        self.ppo_update_counter = 0
+
+    @staticmethod
+    def construct_algorithm(obs: TensorDict, env: VecEnv, cfg: dict, device: str) -> AMPPPO:
+
+        cfg["algorithm"]["style_reward_weight"]= cfg["style_reward_weight"]
+        cfg["algorithm"]["discriminator"] = Discriminator(
+            input_dim_per_frame=cfg["frame_dim"],
+            input_history_length=cfg["frame_length"],
+            hidden_dims=cfg["discriminator_hidden_dims"],
+            feature_normalization=cfg["frame_normalization"],
+            device=device
+        ).to(device)
+
+        cfg["algorithm"]["data_loader"] = AMPDataLoader(
+            cfg["dataset_path"],
+            history_length=cfg["frame_length"],
+            device=device
+        )
+
+        alg:AMPPPO = PPO.construct_algorithm(obs, env, cfg, device)
+
+        return alg
+
+    def process_env_step(
+        self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
+    ) -> None:
+    
+        with torch.no_grad():
+            amp_reward = 0.01*self.style_reward_weight * self.discriminator.get_amp_reward(obs["amp"])
+        
+        task_reward = rewards.clone()
+        total_reward = task_reward + amp_reward
+
+        super().process_env_step(obs, total_reward, dones, extras)
+        
+        # log the single step reward 
+        extras['log']['Step_Reward/style_reward'] = amp_reward
+        extras['log']['Step_Reward/task_reward'] = task_reward
+        
+
+    def update(self):  # noqa: C901
+        loss_dict = {}
+        if self.ppo_update_counter % 4 ==0:
+            mean_amp_policy_score = 0
+            mean_amp_expert_score = 0
+            mean_amp_grad_penalty = 0
+            mean_amp_loss = 0
+            reference_data_generator = self.discriminator_data_loader.mini_batch_generator(self.num_mini_batches//2, self.num_learning_epochs)
+            generator = self.storage.mini_batch_generator(self.num_mini_batches//2, self.num_learning_epochs)
+            # Iterate over batches
+            for (batch,amp_ref_batch) in zip(generator, reference_data_generator):
+
+                amp_policy_batch = batch.observations["amp"]
+                
+                expert_score = self.discriminator(amp_ref_batch)
+                policy_score = self.discriminator(amp_policy_batch)
+
+                expert_loss = torch.nn.MSELoss()(expert_score, torch.ones_like(expert_score))
+                policy_loss = torch.nn.MSELoss()(policy_score, -1 * torch.ones_like(policy_score))
+
+                discrim_loss = 0.5 * (expert_loss + policy_loss)
+
+                grad_pen_loss = self.discriminator.compute_grad_pen(amp_ref_batch)
+
+                discrim_total_loss = discrim_loss + grad_pen_loss
+
+                self.disc_optimizer.zero_grad()
+                discrim_total_loss.backward()
+
+                nn.utils.clip_grad_norm_(self.discriminator.parameters(), self.max_grad_norm)
+                self.disc_optimizer.step()
+        
+                with torch.no_grad():
+                    self.discriminator.update_normalization(amp_ref_batch.detach())
+                    self.discriminator.update_normalization(amp_policy_batch.detach())
+
+                mean_amp_loss += discrim_total_loss.item()
+                mean_amp_grad_penalty += grad_pen_loss.item()
+                mean_amp_expert_score += expert_score.mean().item()
+                mean_amp_policy_score += policy_score.mean().item()
+
+            mean_amp_expert_score /= (self.num_mini_batches * self.num_learning_epochs)
+            mean_amp_policy_score /= (self.num_mini_batches * self.num_learning_epochs)
+            mean_amp_grad_penalty /= (self.num_mini_batches * self.num_learning_epochs)
+            mean_amp_loss /= (self.num_mini_batches * self.num_learning_epochs)
+
+            loss_dict.update({
+                "discriminator_loss": mean_amp_loss,
+                "amp_grad_penalty": mean_amp_grad_penalty,
+                "amp_expert_score": mean_amp_expert_score,  
+                "amp_policy_score": mean_amp_policy_score,
+            })
+            self.policy_update_counter = 0
+            
+
+        ppo_loss_dict = super().update()
+        loss_dict.update(ppo_loss_dict)
+        self.ppo_update_counter += 1
+
+        return loss_dict

source/engineai_lab/assets/__init__.py

@@ -0,0 +1,4 @@
+import os
+
+# Conveniences to other module directories via relative paths
+ASSET_DIR = os.path.abspath(os.path.dirname(__file__))

source/engineai_lab/assets/pm01/urdf/serial_pm01.urdf

@@ -0,0 +1,749 @@
+<?xml version="1.0"?>
+<!-- Use this URDF for robot training instead of the version in the `native_sdk` repository, 
+    due to mismatched inertial properties and collision models. -->
+<robot name="engineai_pm01">
+    <link name="LINK_BASE">
+        <inertial>
+            <origin xyz="0.01525750 -0.00001621 -0.02333111" rpy="0 0 0"/>
+            <mass value="4.08640035"/>
+            <inertia ixx="0.02021593" ixy="-0.00000120" ixz="-0.00037784" iyy="0.01531282" iyz="-0.00000295" izz="0.01216285"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_BASE.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.02 0 -0.042" rpy="0 0 0"/>
+            <geometry>
+                <sphere radius="0.075"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.02 0 0.01" rpy="0 0 0"/>
+            <geometry>
+                <cylinder radius="0.08" length="0.10"/>
+            </geometry>
+        </collision>
+    </link>
+    <link name="LINK_HIP_PITCH_L">
+        <inertial>
+            <origin xyz="0.01010860 0.04524604 -0.01196488" rpy="0 0 0"/>
+            <mass value="1.68628239"/>
+            <inertia ixx="0.00230276" ixy="0.00006587" ixz="-0.00002062" iyy="0.00202247" iyz="-0.00003966" izz="0.00220299"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HIP_PITCH_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0.049359 -0.013226" rpy="0 1.57 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.10"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J00_HIP_PITCH_L" type="revolute">
+        <origin xyz="0.01541 0.076141 -0.061208" rpy="0 0 0"/>
+        <parent link="LINK_BASE"/>
+        <child link="LINK_HIP_PITCH_L"/>
+        <axis xyz="0 0.96593 -0.25882"/>
+        <limit lower="-3.141" upper="2.443" effort="164" velocity="26.3"/>
+    </joint>
+    <link name="LINK_HIP_ROLL_L">
+        <inertial>
+            <origin xyz="-0.01861361 0.00234772 -0.04601977" rpy="0 0 0"/>
+            <mass value="0.63274283"/>
+            <inertia ixx="0.00133858" ixy="-0.00004643" ixz="0.00032023" iyy="0.00144745" iyz="0.00003062" izz="0.00125215"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HIP_ROLL_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="-0.032 0 -0.078" rpy="0 0 0"/>
+            <geometry>
+                <cylinder radius="0.065" length="0.022"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J01_HIP_ROLL_L" type="revolute">
+        <origin xyz="0.048 0.049359 -0.013226" rpy="0 0 0"/>
+        <parent link="LINK_HIP_PITCH_L"/>
+        <child link="LINK_HIP_ROLL_L"/>
+        <axis xyz="1 0 0"/>
+        <limit lower="-0.436" upper="2.094" effort="164" velocity="26.3"/>
+    </joint>
+    <link name="LINK_HIP_YAW_L">
+        <inertial>
+            <origin xyz="-0.00361155 0.00117239 -0.09310066" rpy="0 0 0"/>
+            <mass value="1.85800304"/>
+            <inertia ixx="0.01637894" ixy="-0.00006202" ixz="0.00110380" iyy="0.01626744" iyz="-0.00041393" izz="0.00343388"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HIP_YAW_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 -0.08" rpy="0 0 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.12"/>
+            </geometry>
+        </collision>
+        <!-- <collision>
+            <origin xyz="-0.026 0.038 -0.236" rpy="1.57 0 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.02"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="-0.026 -0.038 -0.236" rpy="1.57 0 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.02"/>
+            </geometry>
+        </collision> -->
+    </link>
+    <joint name="J02_HIP_YAW_L" type="revolute">
+        <origin xyz="-0.03139 -0.0015951 -0.086016" rpy="0 0 0"/>
+        <parent link="LINK_HIP_ROLL_L"/>
+        <child link="LINK_HIP_YAW_L"/>
+        <axis xyz="0 0 1"/>
+        <limit lower="-1.57" upper="4.014" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_KNEE_PITCH_L">
+        <inertial>
+            <origin xyz="-0.00978702 0.00212191 -0.11930917" rpy="0 0 0"/>
+            <mass value="4.29474062"/>
+            <inertia ixx="0.05314257" ixy="0.00007471" ixz="0.00581909" iyy="0.05480624" iyz="0.00057095" izz="0.00561043"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_KNEE_PITCH_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0" rpy="1.57 0 0"/>
+            <geometry>
+                <cylinder radius="0.065" length="0.08"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="-0.02 0 -0.20" rpy="0 0.2 0"/>
+            <geometry>
+                <box size="0.1 0.08 0.15"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J03_KNEE_PITCH_L" type="revolute">
+        <origin xyz="-0.02602 -2.8566E-05 -0.23655" rpy="0 0 0"/>
+        <parent link="LINK_HIP_YAW_L"/>
+        <child link="LINK_KNEE_PITCH_L"/>
+        <axis xyz="0 1 0"/>
+        <limit lower="-0.3491" upper="2.3911" effort="164" velocity="26.3"/>
+    </joint>
+    <link name="LINK_ANKLE_PITCH_L">
+        <inertial>
+            <origin xyz="0.00126412 0 -0.00661445" rpy="0 0 0"/>
+            <mass value="0.11534469"/>
+            <inertia ixx="0.00001346" ixy="0" ixz="0.00000113" iyy="0.00001511" iyz="0" izz="0.00001253"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_PITCH_L.dae"/>
+            </geometry>
+        </visual>
+        <!-- <collision>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_PITCH_L.dae"/>
+            </geometry>
+        </collision> -->
+    </link>
+    <joint name="J04_ANKLE_PITCH_L" type="revolute">
+        <origin xyz="-0.026756 0.00041994 -0.36305" rpy="0 0 0"/>
+        <parent link="LINK_KNEE_PITCH_L"/>
+        <child link="LINK_ANKLE_PITCH_L"/>
+        <axis xyz="0 1 0"/>
+        <limit lower="-0.6807" upper="0.7243" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_ANKLE_ROLL_L">
+        <inertial>
+            <origin xyz="0.01456057 0.00000119 -0.02219893" rpy="0 0 0"/>
+            <mass value="0.71786765"/>
+            <inertia ixx="0.00057796" ixy="-0.00002888" ixz="-0.00014638" iyy="0.00276500" iyz="-0.00000205" izz="0.00292373"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_ROLL_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_ROLL_L.dae"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J05_ANKLE_ROLL_L" type="revolute">
+        <origin xyz="0 0 -0.015" rpy="0 0 0"/>
+        <parent link="LINK_ANKLE_PITCH_L"/>
+        <child link="LINK_ANKLE_ROLL_L"/>
+        <axis xyz="1 0 0"/>
+        <limit lower="-0.2618" upper="0.2618" effort="52" velocity="35.2"/>
+    </joint>
+
+
+    <link name="LINK_HIP_PITCH_R">
+        <inertial>
+            <origin xyz="0.01021848 -0.04516729 -0.01221388" rpy="0 0 0"/>
+            <mass value="1.68024369"/>
+            <inertia ixx="0.00229703" ixy="-0.00006759" ixz="-0.00001661" iyy="0.00202245" iyz="0.00005508" izz="0.00218388"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HIP_PITCH_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 -0.049359 -0.013226" rpy="0 1.57 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.10"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J06_HIP_PITCH_R" type="revolute">
+        <origin xyz="0.01541 -0.076141 -0.061208" rpy="0 0 0"/>
+        <parent link="LINK_BASE"/>
+        <child link="LINK_HIP_PITCH_R"/>
+        <axis xyz="0 0.96593 0.25882"/>
+        <limit lower="-3.141" upper="2.443" effort="164" velocity="26.3"/>
+    </joint>
+    <link name="LINK_HIP_ROLL_R">
+        <inertial>
+            <origin xyz="-0.01847775 -0.00232999 -0.04573211" rpy="0 0 0"/>
+            <mass value="0.63673573"/>
+            <inertia ixx="0.00134870" ixy="0.00004656" ixz="0.00032421" iyy="0.00145899" iyz="-0.00003020" izz="0.00125531"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HIP_ROLL_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="-0.032 0 -0.078" rpy="0 0 0"/>
+            <geometry>
+                <cylinder radius="0.065" length="0.022"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J07_HIP_ROLL_R" type="revolute">
+        <origin xyz="0.048 -0.04936 -0.013226" rpy="0 0 0"/>
+        <parent link="LINK_HIP_PITCH_R"/>
+        <child link="LINK_HIP_ROLL_R"/>
+        <axis xyz="1 0 0"/>
+        <limit lower="-2.094" upper="0.436" effort="164" velocity="26.3"/>
+    </joint>
+    <link name="LINK_HIP_YAW_R">
+        <inertial>
+            <origin xyz="-0.00365713 -0.00305643 -0.09626805" rpy="0 0 0"/>
+            <mass value="1.83146024"/>
+            <inertia ixx="0.01612343" ixy="0.00005892" ixz="0.00108663" iyy="0.01601231" iyz="0.00037741" izz="0.00336783"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HIP_YAW_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 -0.08" rpy="0 0 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.12"/>
+            </geometry>
+        </collision>
+        <!-- <collision>
+            <origin xyz="-0.026 0.038 -0.236" rpy="1.57 0 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.02"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="-0.026 -0.038 -0.236" rpy="1.57 0 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.02"/>
+            </geometry>
+        </collision> -->
+    </link>
+    <joint name="J08_HIP_YAW_R" type="revolute">
+        <origin xyz="-0.03139 0.0015966 -0.086016" rpy="0 0 0"/>
+        <parent link="LINK_HIP_ROLL_R"/>
+        <child link="LINK_HIP_YAW_R"/>
+        <axis xyz="0 0 1"/>
+        <limit lower="-4.014" upper="1.57" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_KNEE_PITCH_R">
+        <inertial>
+            <origin xyz="-0.00985503 -0.00214867 -0.11898555" rpy="0 0 0"/>
+            <mass value="4.29186385"/>
+            <inertia ixx="0.05277030" ixy="-0.00007617" ixz="0.00582343" iyy="0.05446054" iyz="-0.00054823" izz="0.00563558"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_KNEE_PITCH_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0" rpy="1.57 0 0"/>
+            <geometry>
+                <cylinder radius="0.065" length="0.08"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="-0.02 0 -0.20" rpy="0 0.2 0"/>
+            <geometry>
+                <box size="0.1 0.08 0.15"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J09_KNEE_PITCH_R" type="revolute">
+        <origin xyz="-0.02602 2.8566E-05 -0.23655" rpy="0 0 0"/>
+        <parent link="LINK_HIP_YAW_R"/>
+        <child link="LINK_KNEE_PITCH_R"/>
+        <axis xyz="0 1 0"/>
+        <limit lower="-0.3491" upper="2.3911" effort="164" velocity="26.3"/>
+    </joint>
+    <link name="LINK_ANKLE_PITCH_R">
+        <inertial>
+            <origin xyz="0.00073183 0 -0.00669428" rpy="0 0 0"/>
+            <mass value="0.11534469"/>
+            <inertia ixx="0.00001327" ixy="0" ixz="0.00000119" iyy="0.00001511" iyz="0" izz="0.00001272"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_PITCH_R.dae"/>
+            </geometry>
+        </visual>
+        <!-- <collision>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_PITCH_R.dae"/>
+            </geometry>
+        </collision> -->
+    </link>
+    <joint name="J10_ANKLE_PITCH_R" type="revolute">
+        <origin xyz="-0.026756 -0.00041994 -0.36305" rpy="0 0 0"/>
+        <parent link="LINK_KNEE_PITCH_R"/>
+        <child link="LINK_ANKLE_PITCH_R"/>
+        <axis xyz="0 1 0"/>
+        <limit lower="-0.6807" upper="0.7243" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_ANKLE_ROLL_R">
+        <inertial>
+            <origin xyz="0.01470675 0.00001452 -0.02226584" rpy="0 0 0"/>
+            <mass value="0.71559163"/>
+            <inertia ixx="0.00057544" ixy="0.00003022" ixz="-0.00014346" iyy="0.00275723" iyz="0.00000194" izz="0.00291607"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_ROLL_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ANKLE_ROLL_R.dae"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J11_ANKLE_ROLL_R" type="revolute">
+        <origin xyz="0 0 -0.015" rpy="0 0 0"/>
+        <parent link="LINK_ANKLE_PITCH_R"/>
+        <child link="LINK_ANKLE_ROLL_R"/>
+        <axis xyz="1 0 0"/>
+        <limit lower="-0.2618" upper="0.2618" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_TORSO_YAW">
+        <inertial>
+            <origin xyz="-0.01140718 -0.00348592 0.16283761" rpy="0 0 0"/>
+            <mass value="9.01442210"/>
+            <inertia ixx="0.08499169" ixy="-0.00035130" ixz="-0.00487771" iyy="0.06564210" iyz="-0.00043579" izz="0.05757680"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_TORSO_YAW.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.0 0 0.04" rpy="0 0 0"/>
+            <geometry>
+                <box size="0.15 0.17 0.08"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0 0 0.16" rpy="0 0 0"/>
+            <geometry>
+                <box size="0.18 0.2 0.2"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J12_WAIST_YAW" type="revolute">
+        <origin xyz="0.01216 0 0.0809" rpy="0 0 0"/>
+        <parent link="LINK_BASE"/>
+        <child link="LINK_TORSO_YAW"/>
+        <axis xyz="0 0 1"/>
+        <limit lower="-4.014" upper="1.57" effort="52" velocity="35.2"/>
+    </joint>
+
+    <link name="LINK_SHOULDER_PITCH_L">
+        <inertial>
+            <origin xyz="-0.00687390 0.05624677 -0.01564153" rpy="0 0 0"/>
+            <mass value="0.93215458"/>
+            <inertia ixx="0.00123482" ixy="-0.00005768" ixz="0.00002372" iyy="0.00072034" iyz="-0.00018994" izz="0.00100386"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_SHOULDER_PITCH_L.dae"/>
+            </geometry>
+        </visual>
+    </link>
+    <joint name="J13_SHOULDER_PITCH_L" type="revolute">
+        <origin xyz="-0.027105 0.12916 0.21549" rpy="0 0 0"/>
+        <parent link="LINK_TORSO_YAW"/>
+        <child link="LINK_SHOULDER_PITCH_L"/>
+        <axis xyz="0 0.99803 0.062791"/>
+        <limit lower="-2.9671" upper="2.7925" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_SHOULDER_ROLL_L">
+        <inertial>
+            <origin xyz="0.03741198 0.00677175 -0.02678007" rpy="0 0 0"/>
+            <mass value="0.51081280"/>
+            <inertia ixx="0.00099276" ixy="-0.00000156" ixz="0.00000333" iyy="0.00129412" iyz="-0.00012129" izz="0.00093889"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_SHOULDER_ROLL_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.039 0 0" rpy="0 1.57 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.10"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.04 0.02 -0.068" rpy="0.08 0 0"/>
+            <geometry>
+                <cylinder radius="0.04" length="0.032"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J14_SHOULDER_ROLL_L" type="revolute">
+        <origin xyz="-0.0371 0.066941 -0.020838" rpy="0 0 0"/>
+        <parent link="LINK_SHOULDER_PITCH_L"/>
+        <child link="LINK_SHOULDER_ROLL_L"/>
+        <axis xyz="1 0 0"/>
+        <limit lower="-0.6108" upper="2.3562" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_SHOULDER_YAW_L">
+        <inertial>
+            <origin xyz="-0.00084545 0.00219391 -0.04566495" rpy="0 0 0"/>
+            <mass value="0.90913792"/>
+            <inertia ixx="0.00162163" ixy="0.00000177" ixz="0.00000997" iyy="0.00150870" iyz="-0.00001196" izz="0.00078740"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_SHOULDER_YAW_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.0 0.0 -0.038" rpy="0.08 0 0"/>
+            <geometry>
+                <cylinder radius="0.04" length="0.032"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.0 0.008 -0.105" rpy="-1.49 0 0"/>
+            <geometry>
+                <cylinder radius="0.04" length="0.075"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J15_SHOULDER_YAW_L" type="revolute">
+        <origin xyz="0.0371 0.017645 -0.070132" rpy="0 0 0"/>
+        <parent link="LINK_SHOULDER_ROLL_L"/>
+        <child link="LINK_SHOULDER_YAW_L"/>
+        <axis xyz="0 -0.062803 0.99803"/>
+        <limit lower="-2.618" upper="2.618" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_ELBOW_PITCH_L">
+        <inertial>
+            <origin xyz="0.00298237 0.00186711 -0.06518459" rpy="0 0 0"/>
+            <mass value="1.38061447"/>
+            <inertia ixx="0.00588024" ixy="0.00002816" ixz="-0.00030012" iyy="0.00593736" iyz="-0.00046397" izz="0.00083553"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ELBOW_PITCH_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.0 0.007 -0.105" rpy="-1.49 0 0"/>
+            <geometry>
+                <cylinder radius="0.035" length="0.07"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J16_ELBOW_PITCH_L" type="revolute">
+        <origin xyz="0 0.0065994 -0.10487" rpy="0 0 0"/>
+        <parent link="LINK_SHOULDER_YAW_L"/>
+        <child link="LINK_ELBOW_PITCH_L"/>
+        <axis xyz="0.0027243 0.99802 0.062803"/>
+        <limit lower="-2.1948" upper="0.7374" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_ELBOW_YAW_L">
+        <inertial>
+            <origin xyz="0.02016404 0.00026082 -0.08990721" rpy="0 0 0"/>
+            <mass value="0.46651916"/>
+            <inertia ixx="0.00167348" ixy="0.00000913" ixz="0.00030658" iyy="0.00173094" iyz="-0.00004087" izz="0.00039325"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ELBOW_YAW_L.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.01 0.008 -0.05" rpy="0 -0.25 0"/>
+            <geometry>
+                <cylinder radius="0.03" length="0.07"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J17_ELBOW_YAW_L" type="revolute">
+        <origin xyz="0.013817 0.0097723 -0.1547" rpy="0 0 0"/>
+        <parent link="LINK_ELBOW_PITCH_L"/>
+        <child link="LINK_ELBOW_YAW_L"/>
+        <axis xyz="-0.21479 -0.061921 0.9747"/>
+        <limit lower="-2.618" upper="2.618" effort="52" velocity="35.2"/>
+    </joint>
+
+    <link name="LINK_ELBOW_END_L">
+        <inertial>
+            <mass value="1e-3"/>
+            <inertia ixx="1e-12" ixy="0" ixz="0" iyy="1e-12" iyz="0" izz="1e-12"/>
+        </inertial>
+        <collision>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <sphere radius="0.04"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J_FIXED_ELBOW_END_L" type="fixed">
+        <origin xyz="0.03 -0.02 -0.15" rpy="0 0 0"/>
+        <parent link="LINK_ELBOW_YAW_L"/>
+        <child link="LINK_ELBOW_END_L"/>
+    </joint>
+
+    <link name="LINK_SHOULDER_PITCH_R">
+        <inertial>
+            <origin xyz="-0.00682535 -0.05628532 -0.01567159" rpy="0 0 0"/>
+            <mass value="0.92914080"/>
+            <inertia ixx="0.00123547" ixy="0.00005830" ixz="0.00002503" iyy="0.00071785" iyz="0.00018920" izz="0.00100508"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_SHOULDER_PITCH_R.dae"/>
+            </geometry>
+        </visual>
+    </link>
+    <joint name="J18_SHOULDER_PITCH_R" type="revolute">
+        <origin xyz="-0.027105 -0.12916 0.21549" rpy="0 0 0"/>
+        <parent link="LINK_TORSO_YAW"/>
+        <child link="LINK_SHOULDER_PITCH_R"/>
+        <axis xyz="0 0.99803 -0.062791"/>
+        <limit lower="-2.9671" upper="2.7925" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_SHOULDER_ROLL_R">
+        <inertial>
+            <origin xyz="0.03745836 -0.00675724 -0.02671114" rpy="0 0 0"/>
+            <mass value="0.50967809"/>
+            <inertia ixx="0.00099158" ixy="0.00000134" ixz="0.00000304" iyy="0.00129182" iyz="0.00012106" izz="0.00093761"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_SHOULDER_ROLL_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.039 0 0" rpy="0 1.57 0"/>
+            <geometry>
+                <cylinder radius="0.05" length="0.10"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.04 -0.02 -0.068" rpy="-0.08 0 0"/>
+            <geometry>
+                <cylinder radius="0.04" length="0.032"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J19_SHOULDER_ROLL_R" type="revolute">
+        <origin xyz="-0.0371 -0.066941 -0.020838" rpy="0 0 0"/>
+        <parent link="LINK_SHOULDER_PITCH_R"/>
+        <child link="LINK_SHOULDER_ROLL_R"/>
+        <axis xyz="1 0 0"/>
+        <limit lower="-2.3562" upper="0.6108" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_SHOULDER_YAW_R">
+        <inertial>
+            <origin xyz="-0.00081742 -0.00217677 -0.04566906" rpy="0 0 0"/>
+            <mass value="0.90851688"/>
+            <inertia ixx="0.00162102" ixy="-0.00000094" ixz="0.00001055" iyy="0.00150799" iyz="0.00001216" izz="0.00078668"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_SHOULDER_YAW_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.0 0.0 -0.038" rpy="-0.08 0 0"/>
+            <geometry>
+                <cylinder radius="0.04" length="0.032"/>
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.0 -0.008 -0.105" rpy="1.49 0 0"/>
+            <geometry>
+                <cylinder radius="0.04" length="0.075"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J20_SHOULDER_YAW_R" type="revolute">
+        <origin xyz="0.0371 -0.017644 -0.070132" rpy="0 0 0"/>
+        <parent link="LINK_SHOULDER_ROLL_R"/>
+        <child link="LINK_SHOULDER_YAW_R"/>
+        <axis xyz="0 0.062791 0.99803"/>
+        <limit lower="-2.618" upper="2.618" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_ELBOW_PITCH_R">
+        <inertial>
+            <origin xyz="0.00302244 -0.00191447 -0.06521149" rpy="0 0 0"/>
+            <mass value="1.38162913"/>
+            <inertia ixx="0.00587907" ixy="-0.00002756" ixz="-0.00029909" iyy="0.00593696" iyz="0.00046493" izz="0.00083539"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ELBOW_PITCH_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.0 -0.007 -0.105" rpy="1.49 0 0"/>
+            <geometry>
+                <cylinder radius="0.035" length="0.07"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J21_ELBOW_PITCH_R" type="revolute">
+        <origin xyz="0 -0.006598 -0.10487" rpy="0 0 0"/>
+        <parent link="LINK_SHOULDER_YAW_R"/>
+        <child link="LINK_ELBOW_PITCH_R"/>
+        <axis xyz="0.0027243 0.99802 -0.06279"/>
+        <limit lower="-2.1948" upper="0.7374" effort="52" velocity="35.2"/>
+    </joint>
+    <link name="LINK_ELBOW_YAW_R">
+        <inertial>
+            <origin xyz="0.02021657 -0.00026434 -0.08993096" rpy="0 0 0"/>
+            <mass value="0.46648073"/>
+            <inertia ixx="0.00167086" ixy="-0.00000942" ixz="0.00030396" iyy="0.00172726" iyz="0.00004086" izz="0.00039220"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_ELBOW_YAW_R.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0.01 -0.008 -0.05" rpy="0 -0.25 0"/>
+            <geometry>
+                <cylinder radius="0.03" length="0.07"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J22_ELBOW_YAW_R" type="revolute">
+        <origin xyz="0.013817 -0.0097704 -0.1547" rpy="0 0 0"/>
+        <parent link="LINK_ELBOW_PITCH_R"/>
+        <child link="LINK_ELBOW_YAW_R"/>
+        <axis xyz="-0.21479 0.061909 0.9747"/>
+        <limit lower="-2.618" upper="2.618" effort="52" velocity="35.2"/>
+    </joint>
+
+    <link name="LINK_ELBOW_END_R">
+        <inertial>
+            <mass value="1e-3"/>
+            <inertia ixx="1e-12" ixy="0" ixz="0" iyy="1e-12" iyz="0" izz="1e-12"/>
+        </inertial>
+        <collision>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <sphere radius="0.04"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J_FIXED_ELBOW_END_R" type="fixed">
+        <origin xyz="0.03 0.02 -0.15" rpy="0 0 0"/>
+        <parent link="LINK_ELBOW_YAW_R"/>
+        <child link="LINK_ELBOW_END_R"/>
+    </joint>
+
+    <link name="LINK_HEAD_YAW">
+        <inertial>
+            <origin xyz="0.00358213 0.00030109 0.08822633" rpy="0 0 0"/>
+            <mass value="0.84510036"/>
+            <inertia ixx="0.00428261" ixy="0.00000647" ixz="-0.00048256" iyy="0.00500542" iyz="-0.00000033" izz="0.00313801"/>
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0"/>
+            <geometry>
+                <mesh filename="../meshes/LINK_HEAD_YAW.dae"/>
+            </geometry>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0.1" rpy="0 0 0"/>
+            <geometry>
+                <sphere radius="0.08"/>
+            </geometry>
+        </collision>
+    </link>
+    <joint name="J23_HEAD_YAW" type="fixed">
+        <origin xyz="-0.017638 0 0.2961" rpy="0 0 0"/>
+        <parent link="LINK_TORSO_YAW"/>
+        <child link="LINK_HEAD_YAW"/>
+        <axis xyz="0 0 1"/>
+        <limit lower="-0.6109" upper="0.6109" effort="52" velocity="35.2"/>
+    </joint>
+</robot>

source/engineai_lab/robots/actuator.py

@@ -0,0 +1,85 @@
+"""
+This code slice comes from : https://github.com/HybridRobotics/whole_body_tracking
+under the path: source/whole_body_tracking/whole_body_tracking/robots/actuator.py
+"""
+from __future__ import annotations
+
+import torch
+from collections.abc import Sequence
+
+from isaaclab.actuators import ImplicitActuator, ImplicitActuatorCfg
+from isaaclab.utils import DelayBuffer, configclass
+from isaaclab.utils.types import ArticulationActions
+
+
+class DelayedImplicitActuator(ImplicitActuator):
+    """Ideal PD actuator with delayed command application.
+
+    This class extends the :class:`IdealPDActuator` class by adding a delay to the actuator commands. The delay
+    is implemented using a circular buffer that stores the actuator commands for a certain number of physics steps.
+    The most recent actuation value is pushed to the buffer at every physics step, but the final actuation value
+    applied to the simulation is lagged by a certain number of physics steps.
+
+    The amount of time lag is configurable and can be set to a random value between the minimum and maximum time
+    lag bounds at every reset. The minimum and maximum time lag values are set in the configuration instance passed
+    to the class.
+    """
+
+    cfg: DelayedImplicitActuatorCfg
+    """The configuration for the actuator model."""
+
+    def __init__(self, cfg: DelayedImplicitActuatorCfg, *args, **kwargs):
+        super().__init__(cfg, *args, **kwargs)
+        # instantiate the delay buffers
+        self.positions_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
+        self.velocities_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
+        self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
+        # all of the envs
+        self._ALL_INDICES = torch.arange(self._num_envs, dtype=torch.long, device=self._device)
+
+    def reset(self, env_ids: Sequence[int]):
+        super().reset(env_ids)
+        # number of environments (since env_ids can be a slice)
+        if env_ids is None or env_ids == slice(None):
+            num_envs = self._num_envs
+        else:
+            num_envs = len(env_ids)
+        # set a new random delay for environments in env_ids
+        time_lags = torch.randint(
+            low=self.cfg.min_delay,
+            high=self.cfg.max_delay + 1,
+            size=(num_envs,),
+            dtype=torch.int,
+            device=self._device,
+        )
+        # set delays
+        self.positions_delay_buffer.set_time_lag(time_lags, env_ids)
+        self.velocities_delay_buffer.set_time_lag(time_lags, env_ids)
+        self.efforts_delay_buffer.set_time_lag(time_lags, env_ids)
+        # reset buffers
+        self.positions_delay_buffer.reset(env_ids)
+        self.velocities_delay_buffer.reset(env_ids)
+        self.efforts_delay_buffer.reset(env_ids)
+
+    def compute(
+        self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor
+    ) -> ArticulationActions:
+        # apply delay based on the delay the model for all the setpoints
+        control_action.joint_positions = self.positions_delay_buffer.compute(control_action.joint_positions)
+        control_action.joint_velocities = self.velocities_delay_buffer.compute(control_action.joint_velocities)
+        control_action.joint_efforts = self.efforts_delay_buffer.compute(control_action.joint_efforts)
+        # compte actuator model
+        return super().compute(control_action, joint_pos, joint_vel)
+
+
+@configclass
+class DelayedImplicitActuatorCfg(ImplicitActuatorCfg):
+    """Configuration for a delayed PD actuator."""
+
+    class_type: type = DelayedImplicitActuator
+
+    min_delay: int = 0
+    """Minimum number of physics time-steps with which the actuator command may be delayed. Defaults to 0."""
+
+    max_delay: int = 0
+    """Maximum number of physics time-steps with which the actuator command may be delayed. Defaults to 0."""

source/engineai_lab/robots/gen2.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import isaaclab.sim as sim_utils
+from isaaclab.actuators import ImplicitActuatorCfg
+from isaaclab.assets import ArticulationCfg
+from isaaclab.managers import SceneEntityCfg
+
+
+GEN2_ASSET_DIR = Path(__file__).resolve().parents[2] / "gen2_lab" / "assets"
+GEN2_ORIGINAL_URDF_PATH = GEN2_ASSET_DIR / "robot.urdf"
+GEN2_SIMPLIFIED_COLLISION_URDF_PATH = GEN2_ASSET_DIR / "robot_simplified_collision.urdf"
+GEN2_URDF_PATH = (
+    GEN2_SIMPLIFIED_COLLISION_URDF_PATH
+    if GEN2_SIMPLIFIED_COLLISION_URDF_PATH.exists()
+    else GEN2_ORIGINAL_URDF_PATH
+)
+
+GEN2_LEG_JOINT_NAMES = [
+    "left_leg_J1",
+    "left_leg_J2",
+    "left_leg_J3",
+    "left_leg_J4",
+    "left_leg_J5",
+    "left_leg_J6",
+    "right_leg_J1",
+    "right_leg_J2",
+    "right_leg_J3",
+    "right_leg_J4",
+    "right_leg_J5",
+    "right_leg_J6",
+]
+
+GEN2_WAIST_JOINT_NAMES = ["waist_J1", "waist_J2"]
+
+GEN2_ARM_JOINT_NAMES = [
+    "left_arm_J1",
+    "left_arm_J2",
+    "left_arm_J3",
+    "left_arm_J4",
+    "left_arm_J5",
+    "left_arm_J6",
+    "left_arm_J7",
+    "right_arm_J1",
+    "right_arm_J2",
+    "right_arm_J3",
+    "right_arm_J4",
+    "right_arm_J5",
+    "right_arm_J6",
+    "right_arm_J7",
+]
+
+GEN2_DFS_JOINT_NAMES = GEN2_LEG_JOINT_NAMES + GEN2_WAIST_JOINT_NAMES + GEN2_ARM_JOINT_NAMES
+GEN2_DFS_JOINT_ORDER_ASSET_CFG = SceneEntityCfg("robot", joint_names=GEN2_DFS_JOINT_NAMES, preserve_order=True)
+
+GEN2_FEET_BODY_NAMES = ["left_leg_link_6", "right_leg_link_6"]
+
+
+GEN2_CFG = ArticulationCfg(
+    spawn=sim_utils.UrdfFileCfg(
+        asset_path=str(GEN2_URDF_PATH),
+        activate_contact_sensors=True,
+        fix_base=False,
+        replace_cylinders_with_capsules=True,
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=False,
+            retain_accelerations=False,
+            linear_damping=0.0,
+            angular_damping=0.0,
+            max_linear_velocity=1000.0,
+            max_angular_velocity=1000.0,
+            max_depenetration_velocity=1.0,
+        ),
+        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
+            enabled_self_collisions=False,
+            solver_position_iteration_count=8,
+            solver_velocity_iteration_count=4,
+        ),
+        joint_drive=sim_utils.UrdfConverterCfg.JointDriveCfg(
+            gains=sim_utils.UrdfConverterCfg.JointDriveCfg.PDGainsCfg(stiffness=0, damping=0)
+        ),
+    ),
+    init_state=ArticulationCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 1.282),
+        joint_pos={
+            "left_leg_J1": 0.0,
+            "left_leg_J2": 0.0,
+            "left_leg_J3": 0.0,
+            "left_leg_J4": 0.0,
+            "left_leg_J5": 0.0,
+            "left_leg_J6": 0.0,
+            "right_leg_J1": 0.0,
+            "right_leg_J2": 0.0,
+            "right_leg_J3": 0.0,
+            "right_leg_J4": 0.0,
+            "right_leg_J5": 0.0,
+            "right_leg_J6": 0.0,
+            "waist_J1": 0.0,
+            "waist_J2": 0.0,
+            "left_arm_J1": 0.0,
+            "left_arm_J2": 1.4835298641951802,
+            "left_arm_J3": 1.5707963267948966,
+            "left_arm_J4": 0.0,
+            "left_arm_J5": 1.5707963267948966,
+            "left_arm_J6": 0.0,
+            "left_arm_J7": 0.0,
+            "right_arm_J1": 0.0,
+            "right_arm_J2": -1.4835298641951802,
+            "right_arm_J3": -1.5707963267948966,
+            "right_arm_J4": 0.0,
+            "right_arm_J5": 1.5707963267948966,
+            "right_arm_J6": 0.0,
+            "right_arm_J7": 0.0,
+        },
+        joint_vel={".*": 0.0},
+    ),
+    soft_joint_pos_limit_factor=0.9,
+    actuators={
+        "body": ImplicitActuatorCfg(
+            joint_names_expr=[".*"],
+            stiffness={
+                ".*_leg_J1": 180.0,
+                ".*_leg_J2": 160.0,
+                ".*_leg_J3": 90.0,
+                ".*_leg_J4": 220.0,
+                ".*_leg_J5": 55.0,
+                ".*_leg_J6": 55.0,
+                "waist_J.*": 80.0,
+                ".*_arm_J1": 50.0,
+                ".*_arm_J2": 50.0,
+                ".*_arm_J3": 45.0,
+                ".*_arm_J4": 35.0,
+                ".*_arm_J5": 30.0,
+                ".*_arm_J6": 20.0,
+                ".*_arm_J7": 20.0,
+            },
+            damping={
+                ".*_leg_J1": 6.0,
+                ".*_leg_J2": 5.0,
+                ".*_leg_J3": 4.0,
+                ".*_leg_J4": 7.0,
+                ".*_leg_J5": 1.0,
+                ".*_leg_J6": 1.0,
+                "waist_J.*": 4.0,
+                ".*_arm_J1": 0.8,
+                ".*_arm_J2": 0.8,
+                ".*_arm_J3": 0.8,
+                ".*_arm_J4": 0.6,
+                ".*_arm_J5": 0.5,
+                ".*_arm_J6": 0.4,
+                ".*_arm_J7": 0.4,
+            },
+            effort_limit={
+                ".*_leg_J1": 320.0,
+                ".*_leg_J2": 320.0,
+                ".*_leg_J3": 120.0,
+                ".*_leg_J4": 450.0,
+                ".*_leg_J5": 120.0,
+                ".*_leg_J6": 120.0,
+                "waist_J.*": 182.0,
+                ".*_arm_J1": 182.0,
+                ".*_arm_J2": 182.0,
+                ".*_arm_J3": 134.0,
+                ".*_arm_J4": 66.0,
+                ".*_arm_J5": 25.0,
+                ".*_arm_J6": 25.0,
+                ".*_arm_J7": 25.0,
+            },
+            effort_limit_sim={
+                ".*_leg_J1": 320.0,
+                ".*_leg_J2": 320.0,
+                ".*_leg_J3": 120.0,
+                ".*_leg_J4": 450.0,
+                ".*_leg_J5": 120.0,
+                ".*_leg_J6": 120.0,
+                "waist_J.*": 182.0,
+                ".*_arm_J1": 182.0,
+                ".*_arm_J2": 182.0,
+                ".*_arm_J3": 134.0,
+                ".*_arm_J4": 66.0,
+                ".*_arm_J5": 25.0,
+                ".*_arm_J6": 25.0,
+                ".*_arm_J7": 25.0,
+            },
+            velocity_limit={
+                ".*_leg_J1": 26.3,
+                ".*_leg_J2": 26.3,
+                ".*_leg_J3": 35.2,
+                ".*_leg_J4": 26.3,
+                ".*_leg_J5": 35.2,
+                ".*_leg_J6": 35.2,
+                "waist_J.*": 35.2,
+                ".*_arm_J.*": 35.2,
+            },
+            velocity_limit_sim={
+                ".*_leg_J1": 26.3,
+                ".*_leg_J2": 26.3,
+                ".*_leg_J3": 35.2,
+                ".*_leg_J4": 26.3,
+                ".*_leg_J5": 35.2,
+                ".*_leg_J6": 35.2,
+                "waist_J.*": 35.2,
+                ".*_arm_J.*": 35.2,
+            },
+        )
+    },
+)

source/engineai_lab/robots/pm01.py

@@ -0,0 +1,168 @@
+import isaaclab.sim as sim_utils
+from isaaclab.assets import ArticulationCfg
+
+from engineai_lab.assets import ASSET_DIR
+
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.actuators import ImplicitActuatorCfg
+
+
+PM_WAIST_DFS_JOINT_NAMES = [
+    "J00_HIP_PITCH_L",
+    "J01_HIP_ROLL_L",
+    "J02_HIP_YAW_L",
+    "J03_KNEE_PITCH_L",
+    "J04_ANKLE_PITCH_L",
+    "J05_ANKLE_ROLL_L",
+    "J06_HIP_PITCH_R",
+    "J07_HIP_ROLL_R",
+    "J08_HIP_YAW_R",
+    "J09_KNEE_PITCH_R",
+    "J10_ANKLE_PITCH_R",
+    "J11_ANKLE_ROLL_R",
+    "J12_WAIST_YAW",
+    "J13_SHOULDER_PITCH_L",
+    "J14_SHOULDER_ROLL_L",
+    "J15_SHOULDER_YAW_L",
+    "J16_ELBOW_PITCH_L",
+    "J17_ELBOW_YAW_L",
+    "J18_SHOULDER_PITCH_R",
+    "J19_SHOULDER_ROLL_R",
+    "J20_SHOULDER_YAW_R",
+    "J21_ELBOW_PITCH_R",
+    "J22_ELBOW_YAW_R",
+]
+
+PM01_DFS_JOINT_ORDER_ASSET_CFG = SceneEntityCfg("robot", joint_names=PM_WAIST_DFS_JOINT_NAMES, preserve_order=True)
+
+
+PM01_CFG = ArticulationCfg(
+    spawn=sim_utils.UrdfFileCfg(
+        asset_path=f"{ASSET_DIR}/pm01/urdf/serial_pm01.urdf", # full collision
+        activate_contact_sensors=True,
+        fix_base=False,
+        replace_cylinders_with_capsules=True,
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=False,
+            retain_accelerations=False,
+            linear_damping=0.0,
+            angular_damping=0.0,
+            max_linear_velocity=1000.0,
+            max_angular_velocity=1000.0,
+            max_depenetration_velocity=1.0,
+        ),
+        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
+            enabled_self_collisions=True, solver_position_iteration_count=8, solver_velocity_iteration_count=4
+        ),
+        joint_drive=sim_utils.UrdfConverterCfg.JointDriveCfg(
+            gains=sim_utils.UrdfConverterCfg.JointDriveCfg.PDGainsCfg(stiffness=0, damping=0)
+        ),
+    ),
+    init_state=ArticulationCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 0.9),
+        joint_pos={
+            "J00_HIP_PITCH_L": -0.06,
+            "J01_HIP_ROLL_L": 0.0,
+            "J02_HIP_YAW_L": 0.0,
+            "J03_KNEE_PITCH_L": 0.12,
+            "J04_ANKLE_PITCH_L": -0.06,
+            "J05_ANKLE_ROLL_L": 0.0,
+            "J06_HIP_PITCH_R": -0.06,
+            "J07_HIP_ROLL_R": 0.0,
+            "J08_HIP_YAW_R": 0.0,
+            "J09_KNEE_PITCH_R": 0.12,
+            "J10_ANKLE_PITCH_R": -0.06,
+            "J11_ANKLE_ROLL_R": 0.0,
+            "J12_WAIST_YAW": 0.0,
+            "J13_SHOULDER_PITCH_L": 0.0,
+            "J14_SHOULDER_ROLL_L": 0.15,
+            "J15_SHOULDER_YAW_L": 0.0,
+            "J16_ELBOW_PITCH_L": -0.25,
+            "J17_ELBOW_YAW_L": 0.0,
+            "J18_SHOULDER_PITCH_R": 0.0,
+            "J19_SHOULDER_ROLL_R": -0.15,
+            "J20_SHOULDER_YAW_R": 0.0,
+            "J21_ELBOW_PITCH_R": -0.25,
+            "J22_ELBOW_YAW_R": 0.0,
+        },
+        joint_vel={".*": 0.0},
+    ),
+    soft_joint_pos_limit_factor=0.9,
+    actuators={
+        "body": ImplicitActuatorCfg(
+            joint_names_expr=[
+                ".*",
+            ],
+            stiffness={
+                ".*HIP_PITCH.*": 110,
+                ".*HIP_ROLL.*": 70,
+                ".*HIP_YAW.*": 70,
+                ".*KNEE_PITCH.*": 110,
+                ".*ANKLE_PITCH.*": 30,
+                ".*ANKLE_ROLL.*": 30,
+                ".*SHOULDER_PITCH.*": 50,
+                ".*SHOULDER_ROLL.*": 50,
+                ".*SHOULDER_YAW.*": 50,
+                ".*ELBOW_PITCH.*": 50,
+                ".*ELBOW_YAW.*": 50,
+                ".*WAIST_YAW.*": 50,
+            },
+            damping={
+                ".*HIP_PITCH.*": 5.0,
+                ".*HIP_ROLL.*": 3.0,
+                ".*HIP_YAW.*": 3.0,
+                ".*KNEE_PITCH.*": 5.0,
+                ".*ANKLE_PITCH.*": 0.3,
+                ".*ANKLE_ROLL.*": 0.3,
+                ".*SHOULDER_PITCH.*": 0.3,
+                ".*SHOULDER_ROLL.*": 0.3,
+                ".*SHOULDER_YAW.*": 0.3,
+                ".*ELBOW_PITCH.*": 0.3,
+                ".*ELBOW_YAW.*": 0.3,
+                ".*WAIST_YAW.*": 3.0,
+            },
+            effort_limit={
+                ".*HIP_PITCH.*": 164.0,
+                ".*HIP_ROLL.*": 164.0,
+                ".*HIP_YAW.*": 61.0,
+                ".*KNEE_PITCH.*": 164.0,
+                ".*ANKLE_PITCH.*": 54.9,
+                ".*ANKLE_ROLL.*": 54.9,
+                ".*SHOULDER_PITCH.*": 61.0,
+                ".*SHOULDER_ROLL.*": 61.0,
+                ".*SHOULDER_YAW.*": 61.0,
+                ".*ELBOW_PITCH.*": 61.0,
+                ".*ELBOW_YAW.*": 61.0,
+                ".*WAIST_YAW.*": 61.0,
+            },
+            effort_limit_sim={
+                ".*HIP_PITCH.*": 164.0,
+                ".*HIP_ROLL.*": 164.0,
+                ".*HIP_YAW.*": 61.0,
+                ".*KNEE_PITCH.*": 164.0,
+                ".*ANKLE_PITCH.*": 54.9,
+                ".*ANKLE_ROLL.*": 54.9,
+                ".*SHOULDER_PITCH.*": 61.0,
+                ".*SHOULDER_ROLL.*": 61.0,
+                ".*SHOULDER_YAW.*": 61.0,
+                ".*ELBOW_PITCH.*": 61.0,
+                ".*ELBOW_YAW.*": 61.0,
+                ".*WAIST_YAW.*": 61.0,
+            },
+            velocity_limit={
+                ".*HIP_PITCH.*": 26.3,
+                ".*HIP_ROLL.*": 26.3,
+                ".*HIP_YAW.*": 35.2,
+                ".*KNEE_PITCH.*": 26.3,
+                ".*ANKLE_PITCH.*": 35.2,
+                ".*ANKLE_ROLL.*": 35.2,
+                ".*SHOULDER_PITCH.*": 35.2,
+                ".*SHOULDER_ROLL.*": 35.2,
+                ".*SHOULDER_YAW.*": 35.2,
+                ".*ELBOW_PITCH.*": 35.2,
+                ".*ELBOW_YAW.*": 35.2,
+                ".*WAIST_YAW.*": 35.2,
+            },
+        ),
+    },
+)

source/engineai_lab/tasks/__init__.py

@@ -0,0 +1,9 @@
+"""Package containing task implementations for various robotic environments."""
+
+from isaaclab_tasks.utils import import_packages
+from .velocity import *  # noqa
+
+# The blacklist is used to prevent importing configs from sub-packages
+_BLACKLIST_PKGS = ["utils"]
+# Import all configs in this package
+import_packages(__name__, _BLACKLIST_PKGS)

source/engineai_lab/tasks/velocity/config/gen2/__init__.py

@@ -0,0 +1,58 @@
+import gymnasium as gym
+
+from . import agents
+
+
+gym.register(
+    id="Flat-Gen2-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:Gen2FlatEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Gen2FlatPPORunnerCfg",
+    },
+)
+
+
+gym.register(
+    id="Flat-Gen2-Play-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:Gen2FlatEnvCfg_PLAY",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Gen2FlatPPORunnerCfg",
+    },
+)
+
+
+gym.register(
+    id="Flat-Gen2-Speed-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:Gen2SpeedEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Gen2SpeedPPORunnerCfg",
+    },
+)
+
+
+gym.register(
+    id="Flat-Gen2-Natural-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:Gen2NaturalEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Gen2NaturalPPORunnerCfg",
+    },
+)
+
+
+gym.register(
+    id="Flat-Gen2-Natural-Play-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:Gen2NaturalEnvCfg_PLAY",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Gen2NaturalPPORunnerCfg",
+    },
+)

source/engineai_lab/tasks/velocity/config/gen2/agents/__init__.py

@@ -0,0 +1 @@
+

source/engineai_lab/tasks/velocity/config/gen2/agents/rsl_rl_ppo_cfg.py

@@ -0,0 +1,42 @@
+from isaaclab.utils import configclass
+
+from engineai_lab.tasks.velocity.config.pm01.agents.rsl_rl_ppo_cfg import PM01BasePPORunnerCfg
+
+
+@configclass
+class Gen2FlatPPORunnerCfg(PM01BasePPORunnerCfg):
+    max_iterations = 1500
+    experiment_name = "velocity_flat_terrain_gen2"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.actor["hidden_dims"] = [512, 256, 128]
+        self.critic["hidden_dims"] = [512, 256, 128]
+
+
+@configclass
+class Gen2SpeedPPORunnerCfg(Gen2FlatPPORunnerCfg):
+    """Lower-exploration PPO settings for stage-two and stage-three fine-tuning."""
+
+    max_iterations = 400
+
+    def __post_init__(self):
+        super().__post_init__()
+        if isinstance(self.algorithm, dict):
+            self.algorithm["entropy_coef"] = 0.004
+        else:
+            self.algorithm.entropy_coef = 0.004
+
+
+@configclass
+class Gen2NaturalPPORunnerCfg(Gen2SpeedPPORunnerCfg):
+    """Fine-tuning settings with enough exploration to learn coordinated arm swing."""
+
+    max_iterations = 300
+
+    def __post_init__(self):
+        super().__post_init__()
+        if isinstance(self.algorithm, dict):
+            self.algorithm["entropy_coef"] = 0.005
+        else:
+            self.algorithm.entropy_coef = 0.005

source/engineai_lab/tasks/velocity/config/gen2/flat_env_cfg.py

@@ -0,0 +1,822 @@
+from __future__ import annotations
+
+import isaaclab.sim as sim_utils
+from isaaclab.assets import ArticulationCfg, AssetBaseCfg
+from isaaclab.envs import ManagerBasedRLEnvCfg
+from isaaclab.managers import CurriculumTermCfg as CurrTerm
+from isaaclab.managers import EventTermCfg as EventTerm
+from isaaclab.managers import ObservationGroupCfg as ObsGroup
+from isaaclab.managers import ObservationTermCfg as ObsTerm
+from isaaclab.managers import RewardTermCfg as RewTerm
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.managers import TerminationTermCfg as DoneTerm
+from isaaclab.scene import InteractiveSceneCfg
+from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns
+from isaaclab.terrains import TerrainImporterCfg
+from isaaclab.utils import configclass
+from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
+from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
+
+from engineai_lab.robots.actuator import DelayedImplicitActuatorCfg
+from engineai_lab.robots.gen2 import GEN2_CFG, GEN2_DFS_JOINT_NAMES, GEN2_DFS_JOINT_ORDER_ASSET_CFG
+from engineai_lab.robots.gen2 import GEN2_FEET_BODY_NAMES
+from engineai_lab.tasks.velocity import mdp
+from engineai_lab.tasks.velocity.config.pm01.flat_env_cfg import terrain_generator
+
+
+ACTUATOR_DELAY_RANGE = (2, 8)
+GEN2_TORSO_BODY_NAME = "body_link"
+GEN2_PELVIS_BODY_NAME = "waist_link_2"
+GEN2_PELVIS_HEADING_YAW_OFFSET = 1.5708
+GEN2_TORSO_HEIGHT_TARGET = 1.27194
+GEN2_PELVIS_HEIGHT_TARGET = 0.85094
+GEN2_FOOT_SOLE_FRAME_OFFSETS_RPY = (
+    (1.5708, 1.5708, 0.0),
+    (1.5708, 1.5708, 0.0),
+)
+GEN2_DEFAULT_FOOT_POS_PELVIS_FRAME = (
+    (0.096993, 0.120673, -0.785934),
+    (0.093994, -0.120677, -0.785934),
+)
+GEN2_ARM_SWING_JOINT_NAMES = ["left_arm_J1", "right_arm_J1", "left_arm_J4", "right_arm_J4"]
+GEN2_ARM_PHASE_BODY_NAMES = ["left_arm_link_4", "right_arm_link_4"]
+GEN2_WRIST_JOINT_NAMES = [
+    "left_arm_J5",
+    "left_arm_J6",
+    "left_arm_J7",
+    "right_arm_J5",
+    "right_arm_J6",
+    "right_arm_J7",
+]
+
+# Scale PM01-style gait targets by the two robots' approximate kinematic dimensions.
+PM01_REFERENCE_LEG_LENGTH = 0.82
+PM01_REFERENCE_ARM_REACH = 0.575
+GEN2_LEG_LENGTH = abs(GEN2_DEFAULT_FOOT_POS_PELVIS_FRAME[0][2])
+GEN2_ARM_REACH = 0.70
+GEN2_GAIT_PHASE_DISTANCE = 0.25 * GEN2_LEG_LENGTH / PM01_REFERENCE_LEG_LENGTH
+GEN2_SHOULDER_SWING_AMPLITUDE = 0.20 * PM01_REFERENCE_ARM_REACH / GEN2_ARM_REACH
+GEN2_ELBOW_FLEXION = 0.15 * PM01_REFERENCE_ARM_REACH / GEN2_ARM_REACH
+GEN2_ARM_PHASE_AMPLITUDE = 0.10
+
+
+def _build_delayed_actuators():
+    delayed_actuators = {}
+    for name, cfg in GEN2_CFG.actuators.items():
+        delayed_actuators[name] = DelayedImplicitActuatorCfg(
+            joint_names_expr=cfg.joint_names_expr,
+            effort_limit=cfg.effort_limit,
+            effort_limit_sim=cfg.effort_limit_sim,
+            velocity_limit=cfg.velocity_limit,
+            velocity_limit_sim=cfg.velocity_limit_sim,
+            stiffness=cfg.stiffness,
+            damping=cfg.damping,
+            armature=cfg.armature,
+            friction=cfg.friction,
+            dynamic_friction=cfg.dynamic_friction,
+            viscous_friction=cfg.viscous_friction,
+            min_delay=ACTUATOR_DELAY_RANGE[0],
+            max_delay=ACTUATOR_DELAY_RANGE[1],
+        )
+    return delayed_actuators
+
+
+@configclass
+class Gen2SceneCfg(InteractiveSceneCfg):
+    """Configuration for the terrain scene with the Gen2 robot."""
+
+    terrain = TerrainImporterCfg(
+        prim_path="/World/ground",
+        terrain_type="generator",
+        terrain_generator=terrain_generator,
+        max_init_terrain_level=5,
+        collision_group=-1,
+        physics_material=sim_utils.RigidBodyMaterialCfg(
+            friction_combine_mode="multiply",
+            restitution_combine_mode="multiply",
+            static_friction=1.0,
+            dynamic_friction=1.0,
+        ),
+        visual_material=sim_utils.MdlFileCfg(
+            mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/"
+            "TilesMarbleSpiderWhiteBrickBondHoned.mdl",
+            project_uvw=True,
+            texture_scale=(0.25, 0.25),
+        ),
+        debug_vis=False,
+    )
+
+    robot: ArticulationCfg = GEN2_CFG.replace(
+        prim_path="{ENV_REGEX_NS}/Robot",
+        actuators=_build_delayed_actuators(),
+    )
+
+    height_scanner = RayCasterCfg(
+        prim_path="{ENV_REGEX_NS}/Robot/body_link",
+        offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
+        ray_alignment="yaw",
+        pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
+        debug_vis=False,
+        mesh_prim_paths=["/World/ground"],
+    )
+    contact_forces = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True)
+
+    sky_light = AssetBaseCfg(
+        prim_path="/World/skyLight",
+        spawn=sim_utils.DomeLightCfg(
+            intensity=750.0,
+            texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
+        ),
+    )
+
+
+@configclass
+class Gen2Rewards:
+    """Reward terms adapted from PM01, with Gen2 body names and geometry."""
+
+    pelvis_track_lin_vel_xy_exp = RewTerm(
+        func=mdp.track_lin_vel_xy_yaw_frame_exp_body,
+        weight=2.0,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "command_name": "base_velocity",
+            "sigma": 5,
+            "heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+        },
+    )
+    whole_body_track_ang_vel_z_exp = RewTerm(
+        func=mdp.track_ang_vel_z_world_exp_bodies,
+        weight=2.5,
+        params={
+            "asset_cfg": SceneEntityCfg(
+                "robot", body_names=[GEN2_PELVIS_BODY_NAME, GEN2_TORSO_BODY_NAME], preserve_order=True
+            ),
+            "command_name": "base_velocity",
+            "sigma": 5,
+        },
+    )
+    torso_orientation = RewTerm(
+        func=mdp.body_orientation,
+        weight=0.8,
+        params={"asset_cfg": SceneEntityCfg("robot", body_names=GEN2_TORSO_BODY_NAME), "scale": 10.0},
+    )
+    pelvis_height = RewTerm(
+        func=mdp.body_height_tracking,
+        weight=0.4,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "target_height": GEN2_PELVIS_HEIGHT_TARGET,
+            "scale": 15.0,
+        },
+    )
+    torso_height = RewTerm(
+        func=mdp.body_height_tracking,
+        weight=0.1,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_TORSO_BODY_NAME),
+            "target_height": GEN2_TORSO_HEIGHT_TARGET,
+            "scale": 15.0,
+        },
+    )
+    foot_position = RewTerm(
+        func=mdp.feet_position_relative_to_body,
+        weight=0.5,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_FEET_BODY_NAMES),
+            "reference_body_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "command_name": "base_velocity",
+            "stand_threshold": 0.1,
+            "desired_foot_positions": GEN2_DEFAULT_FOOT_POS_PELVIS_FRAME,
+            "heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+        },
+    )
+    feet_orientation = RewTerm(
+        func=mdp.feet_orientation_relative_to_body,
+        weight=0.25,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_FEET_BODY_NAMES),
+            "reference_body_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "command_name": "base_velocity",
+            "stand_threshold": 0.1,
+            "heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+            "foot_frame_offsets_rpy": GEN2_FOOT_SOLE_FRAME_OFFSETS_RPY,
+        },
+    )
+    torso_pelvis_yaw_alignment = RewTerm(
+        func=mdp.body_yaw_alignment,
+        weight=0.3,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_TORSO_BODY_NAME),
+            "reference_body_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "reference_heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+            "scale": 4.0,
+        },
+    )
+    torso_pelvis_yaw_rate = RewTerm(
+        func=mdp.body_yaw_rate_difference_l2,
+        weight=-0.5,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_TORSO_BODY_NAME),
+            "reference_body_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+        },
+    )
+    waist_pos = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.45,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", joint_names=["waist_J1", "waist_J2"]),
+            "scale": 3.0,
+            "tolerance": 0.0,
+        },
+    )
+    waist_vel = RewTerm(
+        func=mdp.joint_vel_l2,
+        weight=-0.02,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["waist_J1", "waist_J2"])},
+    )
+    leg_joint_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={
+            "asset_cfg": SceneEntityCfg(
+                "robot",
+                joint_names=[".*_leg_J2", ".*_leg_J3", ".*_leg_J6"],
+            ),
+            "scale": 3.0,
+        },
+    )
+    arm_primary_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", joint_names=[".*_arm_J1", ".*_arm_J2", ".*_arm_J3", ".*_arm_J4"]),
+            "scale": 3.0,
+        },
+    )
+    arm_distal_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", joint_names=[".*_arm_J5", ".*_arm_J6", ".*_arm_J7"]),
+            "scale": 8.0,
+        },
+    )
+    feet_contact = RewTerm(
+        func=mdp.biped_contact_mode_reward,
+        weight=0.75,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "command_name": "base_velocity",
+            "force_threshold": 5.0,
+            "linear_threshold": 0.1,
+            "angular_threshold": 0.1,
+        },
+    )
+    feet_air_time = RewTerm(
+        func=mdp.feet_air_time_positive_on_contact,
+        weight=2.0,
+        params={
+            "command_name": "base_velocity",
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "min_air_time": 0.05,
+            "max_air_time": 0.25,
+            "linear_threshold": 0.1,
+            "angular_threshold": 0.1,
+        },
+    )
+    feet_air_time_dense = RewTerm(
+        func=mdp.feet_air_time_positive_biped,
+        weight=2.0,
+        params={
+            "command_name": "base_velocity",
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "threshold": 0.25,
+            "linear_threshold": 0.1,
+            "angular_threshold": 0.1,
+        },
+    )
+    swing_foot_clearance = RewTerm(
+        func=mdp.swing_foot_clearance_reward,
+        weight=0.75,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_FEET_BODY_NAMES),
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "command_name": "base_velocity",
+            "target_height": 0.04,
+            "std": 0.05,
+            "force_threshold": 5.0,
+            "linear_threshold": 0.1,
+            "angular_threshold": 0.1,
+        },
+    )
+    foot_stumble = RewTerm(
+        func=mdp.feet_stumble,
+        weight=-1.0,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "tangential_threshold": 2.0,
+            "normal_threshold": 1.0,
+        },
+    )
+    dof_pos_limits = RewTerm(
+        func=mdp.joint_pos_limits,
+        weight=-10.0,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")},
+    )
+    energy_cost = RewTerm(
+        func=mdp.energy_cost_with_curriculum,
+        weight=-0.004,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", joint_names=[".*"]),
+            "start_scale": 0.1,
+            "power": 0.8,
+            "interval_epochs": 200 * 24,
+        },
+    )
+    feet_slide = RewTerm(
+        func=mdp.feet_slide,
+        weight=-0.4,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_FEET_BODY_NAMES),
+        },
+    )
+    dof_vel = RewTerm(
+        func=mdp.joint_vel_l2,
+        weight=-1.0e-5,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])},
+    )
+    dof_acc = RewTerm(
+        func=mdp.joint_acc_l2,
+        weight=-1.25e-8,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])},
+    )
+    action_rate = RewTerm(
+        func=mdp.action_rate_with_curriculum,
+        weight=-0.06,
+        params={"start_scale": 0.1, "power": 0.8, "interval_epochs": 200 * 24},
+    )
+    action_smoothness = RewTerm(
+        func=mdp.action_smoothness_with_curriculum,
+        weight=-0.04,
+        params={"start_scale": 0.1, "power": 0.8, "interval_epochs": 200 * 24},
+    )
+    dof_torque = RewTerm(
+        func=mdp.joint_torques_l2,
+        weight=-1.0e-6,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])},
+    )
+    termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0)
+
+
+@configclass
+class Gen2SpeedRewards(Gen2Rewards):
+    """Additional damping and landing terms for faster command tracking."""
+
+    feet_slide = RewTerm(
+        func=mdp.feet_slide,
+        weight=-0.7,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES),
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_FEET_BODY_NAMES),
+        },
+    )
+
+    pelvis_vertical_velocity = RewTerm(
+        func=mdp.body_vertical_velocity_l2,
+        weight=-2.0,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "deadband": 0.05,
+        },
+    )
+    pelvis_roll_pitch_ang_vel = RewTerm(
+        func=mdp.body_roll_pitch_ang_vel_l2,
+        weight=-0.1,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "deadband": 0.1,
+        },
+    )
+    torso_roll_pitch_ang_vel = RewTerm(
+        func=mdp.body_roll_pitch_ang_vel_l2,
+        weight=-0.1,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_TORSO_BODY_NAME),
+            "deadband": 0.1,
+        },
+    )
+    feet_landing_velocity = RewTerm(
+        func=mdp.feet_landing_velocity,
+        weight=-1.0,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_FEET_BODY_NAMES, preserve_order=True),
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=GEN2_FEET_BODY_NAMES, preserve_order=True),
+            "velocity_threshold": 0.2,
+            "power": 2.0,
+        },
+    )
+
+
+@configclass
+class Gen2NaturalRewards(Gen2SpeedRewards):
+    """Whole-body walking terms with speed-scaled cross-body arm swing."""
+
+    arm_primary_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", joint_names=[".*_arm_J2", ".*_arm_J3"]),
+            "tolerance": 0.05,
+            "scale": 4.0,
+        },
+    )
+    arm_distal_position = RewTerm(
+        func=mdp.joint_deviation_l1_with_deadband,
+        weight=-0.25,
+        params={
+            "asset_cfg": SceneEntityCfg(
+                "robot", joint_names=GEN2_WRIST_JOINT_NAMES, preserve_order=True
+            ),
+            "deadband": 0.03,
+        },
+    )
+    cross_body_arm_swing = RewTerm(
+        func=mdp.cross_body_arm_swing_reward,
+        weight=0.8,
+        params={
+            "arm_asset_cfg": SceneEntityCfg(
+                "robot", joint_names=GEN2_ARM_SWING_JOINT_NAMES, preserve_order=True
+            ),
+            "feet_asset_cfg": SceneEntityCfg(
+                "robot", body_names=GEN2_FEET_BODY_NAMES, preserve_order=True
+            ),
+            "reference_body_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "command_name": "base_velocity",
+            "heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+            "min_forward_speed": 0.1,
+            "full_swing_speed": 0.6,
+            "min_swing_scale": 0.35,
+            "phase_distance": GEN2_GAIT_PHASE_DISTANCE,
+            "shoulder_amplitude": GEN2_SHOULDER_SWING_AMPLITUDE,
+            "elbow_flexion": GEN2_ELBOW_FLEXION,
+            # Gen2's left/right J1 axes are mirrored, so equal joint signs create opposite physical swing.
+            "shoulder_phase_signs": (-1.0, -1.0),
+            "elbow_flexion_signs": (1.0, -1.0),
+            "std": 0.12,
+        },
+    )
+    contralateral_arm_phase = RewTerm(
+        func=mdp.contralateral_arm_phase_reward,
+        weight=0.5,
+        params={
+            "arm_body_cfg": SceneEntityCfg(
+                "robot", body_names=GEN2_ARM_PHASE_BODY_NAMES, preserve_order=True
+            ),
+            "feet_asset_cfg": SceneEntityCfg(
+                "robot", body_names=GEN2_FEET_BODY_NAMES, preserve_order=True
+            ),
+            "reference_body_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+            "command_name": "base_velocity",
+            "heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+            "min_forward_speed": 0.1,
+            "full_swing_speed": 0.6,
+            "min_swing_scale": 0.35,
+            "phase_distance": GEN2_GAIT_PHASE_DISTANCE,
+            "arm_phase_amplitude": GEN2_ARM_PHASE_AMPLITUDE,
+            "std": 0.045,
+        },
+    )
+    arm_swing_velocity = RewTerm(
+        func=mdp.joint_vel_l2,
+        weight=-0.001,
+        params={
+            "asset_cfg": SceneEntityCfg(
+                "robot", joint_names=GEN2_ARM_SWING_JOINT_NAMES, preserve_order=True
+            )
+        },
+    )
+    wrist_velocity = RewTerm(
+        func=mdp.joint_vel_l2,
+        weight=-0.004,
+        params={
+            "asset_cfg": SceneEntityCfg(
+                "robot", joint_names=GEN2_WRIST_JOINT_NAMES, preserve_order=True
+            )
+        },
+    )
+
+
+@configclass
+class Gen2Termination:
+    time_out = DoneTerm(func=mdp.time_out, time_out=True)
+    base_contact = DoneTerm(
+        func=mdp.illegal_contact,
+        params={
+            "sensor_cfg": SceneEntityCfg(
+                "contact_forces",
+                body_names=["body_link", "waist_link_.*", ".*_leg_link_[1-4]", ".*_arm_link_.*"],
+            ),
+            "threshold": 1.0,
+        },
+    )
+
+
+@configclass
+class ActionsCfg:
+    joint_pos = mdp.JointPositionActionCfg(
+        asset_name="robot",
+        use_default_offset=True,
+        preserve_order=True,
+        joint_names=GEN2_DFS_JOINT_NAMES,
+        scale={
+            ".*_leg_J1": 0.5,
+            ".*_leg_J2": 0.2,
+            ".*_leg_J3": 0.2,
+            ".*_leg_J4": 0.5,
+            ".*_leg_J5": 0.5,
+            ".*_leg_J6": 0.2,
+            "waist_J.*": 0.2,
+            ".*_arm_J1": 0.2,
+            ".*_arm_J2": 0.2,
+            ".*_arm_J3": 0.2,
+            ".*_arm_J4": 0.2,
+            ".*_arm_J5": 0.15,
+            ".*_arm_J6": 0.15,
+            ".*_arm_J7": 0.15,
+        },
+    )
+
+
+@configclass
+class ObservationsCfg:
+    """Observation specifications for the MDP."""
+
+    @configclass
+    class PolicyCfg(ObsGroup):
+        joint_pos = ObsTerm(
+            func=mdp.joint_pos_rel,
+            noise=Unoise(n_min=-0.01, n_max=0.01),
+            params={"asset_cfg": GEN2_DFS_JOINT_ORDER_ASSET_CFG},
+            history_length=15,
+        )
+        joint_vel = ObsTerm(
+            func=mdp.joint_vel_rel,
+            noise=Unoise(n_min=-1.5, n_max=1.5),
+            params={"asset_cfg": GEN2_DFS_JOINT_ORDER_ASSET_CFG},
+            history_length=15,
+        )
+        actions = ObsTerm(func=mdp.last_action, history_length=15)
+        pelvis_ang_vel = ObsTerm(
+            func=mdp.body_ang_vel_yaw_frame,
+            noise=Unoise(n_min=-0.2, n_max=0.2),
+            params={
+                "asset_cfg": SceneEntityCfg("robot", body_names=GEN2_PELVIS_BODY_NAME),
+                "heading_yaw_offset": GEN2_PELVIS_HEADING_YAW_OFFSET,
+            },
+            history_length=15,
+        )
+        torso_projected_gravity = ObsTerm(
+            func=mdp.body_projected_gravity,
+            noise=Unoise(n_min=-0.05, n_max=0.05),
+            params={"asset_cfg": SceneEntityCfg("robot", body_names=GEN2_TORSO_BODY_NAME)},
+            history_length=15,
+        )
+        velocity_commands = ObsTerm(func=mdp.generated_commands, params={"command_name": "base_velocity"})
+
+        def __post_init__(self):
+            self.enable_corruption = True
+            self.concatenate_terms = True
+
+    @configclass
+    class CriticCfg(PolicyCfg):
+        pass
+
+    policy: PolicyCfg = PolicyCfg()
+    critic: CriticCfg = CriticCfg()
+
+
+@configclass
+class Gen2Commands:
+    """Conservative velocity commands for the first Gen2 walking stage."""
+
+    base_velocity = mdp.BodyVelocityCommandCfg(
+        asset_name="robot",
+        body_name=GEN2_PELVIS_BODY_NAME,
+        heading_yaw_offset=GEN2_PELVIS_HEADING_YAW_OFFSET,
+        resampling_time_range=(7.5, 7.5),
+        rel_standing_envs=0.2,
+        rel_heading_envs=0.0,
+        heading_command=False,
+        heading_control_stiffness=0.5,
+        debug_vis=False,
+        ranges=mdp.BodyVelocityCommandCfg.Ranges(
+            lin_vel_x=(-0.2, 0.4),
+            lin_vel_y=(-0.2, 0.2),
+            ang_vel_z=(-0.5, 0.5),
+            heading=None,
+        ),
+    )
+
+
+VELOCITY_RANGE = {
+    "x": (-0.3, 0.3),
+    "y": (-0.3, 0.3),
+    "z": (-0.15, 0.15),
+    "roll": (-0.35, 0.35),
+    "pitch": (-0.35, 0.35),
+    "yaw": (-0.52, 0.52),
+}
+
+
+@configclass
+class Gen2EventCfg:
+    """Gen2-specific randomizations."""
+
+    physics_material = EventTerm(
+        func=mdp.randomize_rigid_body_material,
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
+            "static_friction_range": (0.3, 1.6),
+            "dynamic_friction_range": (0.3, 1.2),
+            "restitution_range": (0.0, 0.5),
+            "num_buckets": 64,
+        },
+    )
+    add_joint_default_pos = EventTerm(
+        func=mdp.randomize_joint_default_pos,
+        mode="startup",
+        params={
+            "asset_cfg": GEN2_DFS_JOINT_ORDER_ASSET_CFG,
+            "pos_distribution_params": (-0.01, 0.01),
+            "operation": "add",
+        },
+    )
+    base_com = EventTerm(
+        func=mdp.randomize_rigid_body_com,
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names="body_link"),
+            "com_range": {"x": (-0.03, 0.03), "y": (-0.06, 0.06), "z": (-0.06, 0.06)},
+        },
+    )
+    push_robot = EventTerm(
+        func=mdp.push_by_setting_velocity,
+        mode="interval",
+        interval_range_s=(1.0, 3.0),
+        params={"velocity_range": VELOCITY_RANGE},
+    )
+    reset_base = EventTerm(
+        func=mdp.reset_root_state_uniform,
+        mode="reset",
+        params={
+            "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)},
+            "velocity_range": {
+                "x": (0.0, 0.0),
+                "y": (0.0, 0.0),
+                "z": (0.0, 0.0),
+                "roll": (0.0, 0.0),
+                "pitch": (0.0, 0.0),
+                "yaw": (0.0, 0.0),
+            },
+        },
+    )
+    reset_robot_joints = EventTerm(
+        func=mdp.reset_joints_by_scale,
+        mode="reset",
+        params={"position_range": (0.8, 1.2), "velocity_range": (-0.5, 0.5)},
+    )
+
+
+@configclass
+class CurriculumCfg:
+    terrain_levels = CurrTerm(func=mdp.terrain_levels_vel)
+
+
+@configclass
+class Gen2FlatEnvCfg(ManagerBasedRLEnvCfg):
+    """Velocity-tracking RL environment for Gen2."""
+
+    scene: Gen2SceneCfg = Gen2SceneCfg(num_envs=4096, env_spacing=3.0)
+    observations: ObservationsCfg = ObservationsCfg()
+    actions: ActionsCfg = ActionsCfg()
+    commands: Gen2Commands = Gen2Commands()
+    rewards: Gen2Rewards = Gen2Rewards()
+    terminations: Gen2Termination = Gen2Termination()
+    events: Gen2EventCfg = Gen2EventCfg()
+    curriculum: CurriculumCfg = CurriculumCfg()
+
+    def __post_init__(self):
+        self.decimation = 5
+        self.episode_length_s = 20.0
+        self.sim.dt = 0.002
+        self.sim.render_interval = self.decimation
+        self.sim.physics_material = self.scene.terrain.physics_material
+        self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15
+        if self.scene.height_scanner is not None:
+            self.scene.height_scanner.update_period = self.decimation * self.sim.dt
+        if self.scene.contact_forces is not None:
+            self.scene.contact_forces.update_period = 0.005
+
+        if getattr(self.curriculum, "terrain_levels", None) is not None:
+            if self.scene.terrain.terrain_generator is not None:
+                self.scene.terrain.terrain_generator.curriculum = True
+        else:
+            if self.scene.terrain.terrain_generator is not None:
+                self.scene.terrain.terrain_generator.curriculum = False
+
+
+@configclass
+class Gen2SpeedEnvCfg(Gen2FlatEnvCfg):
+    """Stage-two flat-ground consolidation task for faster and steadier tracking."""
+
+    rewards: Gen2SpeedRewards = Gen2SpeedRewards()
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        self.commands.base_velocity.resampling_time_range = (3.0, 5.0)
+        self.commands.base_velocity.rel_standing_envs = 0.1
+        self.commands.base_velocity.rel_straight_envs = 0.8
+        self.commands.base_velocity.ranges.lin_vel_x = (0.2, 0.6)
+        self.commands.base_velocity.ranges.lin_vel_y = (-0.05, 0.05)
+        self.commands.base_velocity.ranges.ang_vel_z = (-0.2, 0.2)
+        self.rewards.pelvis_track_lin_vel_xy_exp.weight = 2.5
+
+        self.scene.terrain.terrain_type = "plane"
+        self.scene.terrain.terrain_generator = None
+        self.scene.terrain.max_init_terrain_level = None
+        self.scene.height_scanner = None
+        self.curriculum.terrain_levels = None
+        self.events.push_robot = None
+
+
+@configclass
+class Gen2NaturalEnvCfg(Gen2SpeedEnvCfg):
+    """Stage-three fine-tuning task that adds natural cross-body arm swing."""
+
+    rewards: Gen2NaturalRewards = Gen2NaturalRewards()
+
+    def __post_init__(self):
+        super().__post_init__()
+        for joint_pattern in (".*_arm_J5", ".*_arm_J6", ".*_arm_J7"):
+            self.actions.joint_pos.scale[joint_pattern] = 0.06
+        self.events.reset_robot_joints.params["position_range"] = (0.95, 1.05)
+        self.events.reset_robot_joints.params["velocity_range"] = (-0.2, 0.2)
+
+
+def _configure_gen2_play_env(cfg):
+    """Apply deterministic replay settings without changing a task's action mapping."""
+    cfg.seed = 42
+    cfg.scene.num_envs = 1
+    cfg.scene.env_spacing = 2.5
+    cfg.episode_length_s = 40.0
+
+    cfg.scene.terrain.terrain_type = "plane"
+    cfg.scene.terrain.terrain_generator = None
+    cfg.scene.terrain.max_init_terrain_level = None
+    cfg.scene.height_scanner = None
+    cfg.curriculum.terrain_levels = None
+
+    cfg.observations.policy.enable_corruption = False
+    cfg.observations.critic.enable_corruption = False
+
+    cfg.events.physics_material.params["static_friction_range"] = (1.0, 1.0)
+    cfg.events.physics_material.params["dynamic_friction_range"] = (1.0, 1.0)
+    cfg.events.physics_material.params["restitution_range"] = (0.0, 0.0)
+    cfg.events.add_joint_default_pos = None
+    cfg.events.base_com = None
+    cfg.events.push_robot = None
+    cfg.events.reset_base.params["pose_range"] = {
+        "x": (0.0, 0.0),
+        "y": (0.0, 0.0),
+        "yaw": (0.0, 0.0),
+    }
+    cfg.events.reset_robot_joints.params["position_range"] = (1.0, 1.0)
+    cfg.events.reset_robot_joints.params["velocity_range"] = (0.0, 0.0)
+
+    for actuator_cfg in cfg.scene.robot.actuators.values():
+        if hasattr(actuator_cfg, "min_delay") and hasattr(actuator_cfg, "max_delay"):
+            actuator_cfg.min_delay = 5
+            actuator_cfg.max_delay = 5
+
+    cfg.commands.base_velocity.debug_vis = True
+
+
+@configclass
+class Gen2FlatEnvCfg_PLAY(Gen2FlatEnvCfg):
+    """Deterministic flat-ground configuration for policy replay."""
+
+    def __post_init__(self):
+        super().__post_init__()
+        _configure_gen2_play_env(self)
+
+
+@configclass
+class Gen2NaturalEnvCfg_PLAY(Gen2NaturalEnvCfg):
+    """Deterministic replay configuration preserving stage-three wrist action scales."""
+
+    def __post_init__(self):
+        super().__post_init__()
+        _configure_gen2_play_env(self)

source/engineai_lab/tasks/velocity/config/pm01/__init__.py

@@ -0,0 +1,28 @@
+import gymnasium as gym
+
+from . import agents
+
+##
+# Register Gym environments.
+##
+
+gym.register(
+    id="Flat-PM01-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:PM01FlatEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:PM01FlatPPORunnerCfg",
+    },
+)
+
+
+gym.register(
+    id="Flat-AMP-PM01-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.flat_amp_env_cfg:PM01AMPFlatEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.amp_ppo_cfg:PM01FlatAMPPPORunnerCfg",
+    },
+)

source/engineai_lab/tasks/velocity/config/pm01/agents/amp_ppo_cfg.py

@@ -0,0 +1,22 @@
+from isaaclab.utils import configclass
+
+from .rsl_rl_ppo_cfg import PM01BasePPORunnerCfg
+
+
+@configclass
+class PM01FlatAMPPPORunnerCfg(PM01BasePPORunnerCfg):
+    max_iterations: int = 80_000
+    save_interval: int = 500
+    # AMP parameters
+    style_reward_weight = 2.0
+    frame_length = 5
+    frame_dim = 26
+    frame_normalization = True
+    discriminator_hidden_dims = [512, 256, 128]
+    #
+    experiment_name = "velocity_flat_terrain_amp"
+    dataset_path = "dataset/config/dataset.yaml"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.algorithm.class_name="engineai_lab.algorithms.amp_ppo:AMPPPO"

source/engineai_lab/tasks/velocity/config/pm01/agents/rsl_rl_ppo_cfg.py

@@ -0,0 +1,66 @@
+from isaaclab.utils import configclass
+
+from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg,RslRlMLPModelCfg,RslRlPpoAlgorithmCfg
+
+
+@configclass
+class PM01BasePPORunnerCfg(RslRlOnPolicyRunnerCfg):
+    num_steps_per_env = 24
+    max_iterations = 3000
+    save_interval = 50
+    experiment_name = "velocity_flat_terrain"
+    obs_groups = {"actor": ["policy"], "critic": ["policy"]}
+    algorithm = RslRlPpoAlgorithmCfg(
+        value_loss_coef=1.0,
+        use_clipped_value_loss=True,
+        clip_param=0.2,
+        entropy_coef=0.008,
+        num_learning_epochs=5,
+        num_mini_batches=4,
+        learning_rate=1.0e-3,
+        schedule="adaptive",
+        gamma=0.99,
+        lam=0.95,
+        desired_kl=0.01,
+        max_grad_norm=1.0,
+    )
+    actor = RslRlMLPModelCfg(
+        hidden_dims=[512, 256, 128],
+        activation="elu",
+        obs_normalization=True,
+        distribution_cfg=
+            RslRlMLPModelCfg.GaussianDistributionCfg(
+                init_std=1.0,
+                std_type="scalar"
+            )
+    )
+    critic = RslRlMLPModelCfg(
+        hidden_dims=[512, 256, 128],
+        activation="elu",
+        obs_normalization=True,
+    )
+
+
+    def __post_init__(self):
+        super().__post_init__()
+        
+        # 2. deprecated key words for rsl-rl >= 5.0.0  
+        deprecated_keys = {"stochastic", "init_noise_std", "noise_std_type", "state_dependent_std"}
+        
+        def _remove_deprecated_keys(cfg_obj):
+            if cfg_obj is None:
+                return None
+            return {k: v for k, v in vars(cfg_obj).items() if k not in deprecated_keys}
+        
+        self.actor = _remove_deprecated_keys(self.actor)
+        self.critic = _remove_deprecated_keys(self.critic)
+
+@configclass
+class PM01FlatPPORunnerCfg(PM01BasePPORunnerCfg):
+    max_iterations = 1500
+    experiment_name = "velocity_flat_terrain"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.policy.actor_hidden_dims = [128, 128, 128]
+        self.policy.critic_hidden_dims = [128, 128, 128]

source/engineai_lab/tasks/velocity/config/pm01/flat_amp_env_cfg.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import torch
+
+
+
+from isaaclab.managers import ObservationGroupCfg as ObsGroup
+from isaaclab.managers import ObservationTermCfg as ObsTerm
+from isaaclab.managers import SceneEntityCfg
+
+# Pre-defined configs
+##
+from isaaclab.utils import configclass
+from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
+
+from engineai_lab.tasks.velocity import mdp
+from .flat_env_cfg import PM01FlatEnvCfg
+from engineai_lab.robots.pm01 import PM01_CFG, PM_WAIST_DFS_JOINT_NAMES, PM01_DFS_JOINT_ORDER_ASSET_CFG
+
+def dummy_history_term(env):
+    # zero-out waist joint observations to avoid AMP mismatch between robot variants
+    joint_pos = mdp.joint_pos(env).clone()
+    joint_vel = mdp.joint_vel(env).clone()
+    waist_joint_ids = env.scene["robot"].find_joints(".*WAIST_YAW.*", preserve_order=True)[0]
+    if len(waist_joint_ids) > 0:
+        waist_joint_ids = torch.as_tensor(waist_joint_ids, device=joint_pos.device)
+        joint_pos[:, waist_joint_ids] = 0.0
+        joint_vel[:, waist_joint_ids] = 0.0
+    lin_vel = mdp.robot_base_lin_vel_b(env)
+
+    return torch.cat([joint_pos * 9, lin_vel * 7], dim=-1)
+
+
+@configclass
+class ObservationsCfg:
+    """Observation specifications for the MDP."""
+
+    @configclass
+    class PolicyCfg(ObsGroup):
+        """Observations for policy group."""
+
+        # observation terms (order preserved)
+        joint_pos = ObsTerm(
+            func=mdp.joint_pos_rel,
+            noise=Unoise(n_min=-0.01, n_max=0.01),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        joint_vel = ObsTerm(
+            func=mdp.joint_vel_rel,
+            noise=Unoise(n_min=-1.5, n_max=1.5),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        actions = ObsTerm(func=mdp.last_action,
+                          history_length=15)        
+        base_ang_vel = ObsTerm(func=mdp.base_ang_vel, 
+                               noise=Unoise(n_min=-0.2, n_max=0.2),
+                            history_length=15)
+        projected_gravity = ObsTerm(
+            func=mdp.projected_gravity,
+            noise=Unoise(n_min=-0.05, n_max=0.05),
+            history_length=15
+        )
+        velocity_commands = ObsTerm(func=mdp.generated_commands, 
+                                    params={"command_name": "base_velocity"})        
+        
+        def __post_init__(self):
+            self.enable_corruption = True
+            self.concatenate_terms = True
+    
+    @configclass
+    class CriticCfg(ObsGroup):
+
+        # observation terms (order preserved)
+        joint_pos = ObsTerm(
+            func=mdp.joint_pos_rel,
+            noise=Unoise(n_min=-0.01, n_max=0.01),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        joint_vel = ObsTerm(
+            func=mdp.joint_vel_rel,
+            noise=Unoise(n_min=-1.5, n_max=1.5),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        actions = ObsTerm(func=mdp.last_action,
+                          history_length=15)        
+        base_ang_vel = ObsTerm(func=mdp.base_ang_vel, 
+                               noise=Unoise(n_min=-0.2, n_max=0.2),
+                            history_length=15)
+        projected_gravity = ObsTerm(
+            func=mdp.projected_gravity,
+            noise=Unoise(n_min=-0.05, n_max=0.05),
+            history_length=15
+        )
+        velocity_commands = ObsTerm(func=mdp.generated_commands, 
+                                    params={"command_name": "base_velocity"})
+
+
+        def __post_init__(self):
+            self.enable_corruption = True
+            self.concatenate_terms = True
+
+    @configclass
+    class AMPCfg(ObsGroup):
+        history = ObsTerm(func=dummy_history_term)
+        def __post_init__(self):
+            self.history_length = 5 # TODO: is the history from old to new or new to old?
+        
+
+    # observation groups
+    policy: PolicyCfg = PolicyCfg()
+    critic: CriticCfg = CriticCfg()
+    amp: AMPCfg = AMPCfg()
+    
+
+@configclass
+class PM01AMPFlatEnvCfg(PM01FlatEnvCfg):
+    observations: ObservationsCfg = ObservationsCfg()
+
+    def __post_init__(self):
+        # post init of parent
+        super().__post_init__()
+        

source/engineai_lab/tasks/velocity/config/pm01/flat_env_cfg.py

@@ -0,0 +1,615 @@
+from __future__ import annotations
+import isaaclab.sim as sim_utils
+from isaaclab.assets import ArticulationCfg, AssetBaseCfg
+from isaaclab.envs import ManagerBasedRLEnvCfg
+from isaaclab.managers import CurriculumTermCfg as CurrTerm
+from isaaclab.managers import ObservationGroupCfg as ObsGroup
+from isaaclab.managers import ObservationTermCfg as ObsTerm
+from isaaclab.managers import RewardTermCfg as RewTerm
+from isaaclab.managers import SceneEntityCfg 
+from isaaclab.managers import TerminationTermCfg as DoneTerm
+from isaaclab.managers import EventTermCfg as EventTerm
+from isaaclab.scene import InteractiveSceneCfg
+from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns
+from isaaclab.terrains import TerrainImporterCfg
+##
+# Pre-defined configs
+##
+from isaaclab.utils import configclass
+from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR, ISAAC_NUCLEUS_DIR
+from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
+
+from engineai_lab.tasks.velocity import mdp
+from engineai_lab.robots.pm01 import PM01_CFG, PM_WAIST_DFS_JOINT_NAMES, PM01_DFS_JOINT_ORDER_ASSET_CFG
+
+import isaaclab.terrains as terrain_gen
+import math
+from isaaclab.terrains.terrain_generator_cfg import TerrainGeneratorCfg
+
+from engineai_lab.robots.actuator import DelayedImplicitActuatorCfg
+
+
+ACTUATOR_DELAY_RANGE = (2, 8)
+def _build_delayed_actuators():
+    delayed_actuators = {}
+    for name, cfg in PM01_CFG.actuators.items():
+        delayed_actuators[name] = DelayedImplicitActuatorCfg(
+            joint_names_expr=cfg.joint_names_expr,
+            effort_limit=cfg.effort_limit,
+            effort_limit_sim=cfg.effort_limit_sim,
+            velocity_limit=cfg.velocity_limit,
+            velocity_limit_sim=cfg.velocity_limit_sim,
+            stiffness=cfg.stiffness,
+            damping=cfg.damping,
+            armature=cfg.armature,
+            friction=cfg.friction,
+            dynamic_friction=cfg.dynamic_friction,
+            viscous_friction=cfg.viscous_friction,
+            min_delay=ACTUATOR_DELAY_RANGE[0],
+            max_delay=ACTUATOR_DELAY_RANGE[1],
+        )
+    return delayed_actuators
+
+
+terrain_generator=TerrainGeneratorCfg(
+            size=(8.0, 8.0),
+            horizontal_scale=0.1,
+            vertical_scale=0.005,
+            border_width=25.0,
+            num_rows=10,
+            num_cols=20,
+            curriculum=True,
+            difficulty_range=(0.0, 1.0),
+            color_scheme="height",
+            slope_threshold=0.75,
+            sub_terrains={
+                "flat": terrain_gen.HfPyramidSlopedTerrainCfg(
+                    proportion=0.4,
+                    slope_range=(0.0, 0.0),
+                    platform_width=8.0,
+                ),
+                "slope_up": terrain_gen.HfPyramidSlopedTerrainCfg(
+                    proportion=0.1,
+                    slope_range=(0.0, math.radians(5)),
+                    platform_width=2.0,
+                ),
+                "slope_down": terrain_gen.HfInvertedPyramidSlopedTerrainCfg(
+                    proportion=0.1,
+                    slope_range=(0.0, math.radians(5)),
+                    platform_width=2.0,
+                ),
+                "obstacles": terrain_gen.HfDiscreteObstaclesTerrainCfg(
+                    proportion=0.2,
+                    obstacle_width_range=(1.0, 2.0),
+                    obstacle_height_range=(0.01, 0.1),
+                    num_obstacles=15,
+                    platform_width=3.0,
+                ),
+                "rough_terrain": terrain_gen.HfRandomUniformTerrainCfg(
+                    proportion=0.2,
+                    noise_range=(-0.015, 0.015),
+                    noise_step=0.005,
+                    downsampled_scale=0.15,
+                ),
+            },
+        )
+
+@configclass
+class PM01SceneCfg(InteractiveSceneCfg):
+    """Configuration for the terrain scene with a legged robot."""
+
+    # ground terrain
+    terrain = TerrainImporterCfg(
+        prim_path="/World/ground",
+        terrain_type="generator",
+        terrain_generator=terrain_generator,
+        max_init_terrain_level=5,
+        collision_group=-1,
+        physics_material=sim_utils.RigidBodyMaterialCfg(
+            friction_combine_mode="multiply",
+            restitution_combine_mode="multiply",
+            static_friction=1.0,
+            dynamic_friction=1.0,
+        ),
+        visual_material=sim_utils.MdlFileCfg(
+            mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl",
+            project_uvw=True,
+            texture_scale=(0.25, 0.25),
+        ),
+        debug_vis=False,
+    )
+    # robots
+    robot: ArticulationCfg = PM01_CFG.replace(
+        prim_path="{ENV_REGEX_NS}/Robot",
+        actuators=_build_delayed_actuators(),
+    )
+    # sensors
+    height_scanner = RayCasterCfg(
+        prim_path="{ENV_REGEX_NS}/Robot/LINK_BASE",
+        offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
+        ray_alignment="yaw",
+        pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
+        debug_vis=False,
+        mesh_prim_paths=["/World/ground"],
+    )
+    contact_forces = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True)
+    # lights
+    sky_light = AssetBaseCfg(
+        prim_path="/World/skyLight",
+        spawn=sim_utils.DomeLightCfg(
+            intensity=750.0,
+            texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
+        ),
+    )
+
+
+
+@configclass
+class PM01Rewards:
+    """Reward terms for the MDP."""
+
+
+    track_lin_vel_xy_exp = RewTerm(
+        func=mdp.track_lin_vel_xy_yaw_frame_exp,
+        weight=2.0,
+        params={"command_name": "base_velocity", "sigma": 5},
+    )
+    track_ang_vel_z_exp = RewTerm(
+        func=mdp.track_ang_vel_z_world_exp, 
+        weight=2.5, 
+        params={"command_name": "base_velocity", "sigma": 5}
+    )
+
+    base_orientation = RewTerm(
+        func=mdp.base_orientation,
+        weight=1.0
+    )
+
+    base_height = RewTerm(
+        func=mdp.base_height_tracking,
+        weight=0.4,
+        params={"target_height": 0.82}
+    )
+    
+    foot_position = RewTerm(
+        func=mdp.feet_position,
+        weight=0.5,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "command_name": "base_velocity",
+            "stand_threshold": 0.1,
+            "ankle_distance": 0.22,
+            "base_height_target": 0.82,
+        },
+    )
+    
+    feet_orientation = RewTerm(
+        func=mdp.feet_orientation,
+        weight=0.25,
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "command_name": "base_velocity",
+            "stand_threshold": 0.1,
+        },
+    )
+    
+    waist_pos = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={
+            "asset_cfg":  SceneEntityCfg("robot", joint_names=["J12_WAIST_YAW"]),
+            "scale": 3.0,
+            "tolerance": 0.0
+        },
+    )
+    
+    leg_joint_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_HIP_ROLL_.*", ".*_HIP_YAW_.*", ".*_ANKLE_ROLL_.*"]),
+                "scale": 3.0},
+    )
+    
+    arm_pitch_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*SHOULDER_PITCH.*", ".*ELBOW_PITCH.*"]),
+                "scale": 3.0},
+    )
+    
+    arm_roll_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*SHOULDER_ROLL.*"]),
+                "scale": 3.0},
+    )
+    
+    arm_yaw_position = RewTerm(
+        func=mdp.joint_deviation_exp,
+        weight=0.3,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*SHOULDER_YAW.*", ".*ELBOW_YAW.*"]),
+                "scale": 10.0},
+    )
+    
+    feet_contact = RewTerm(
+        func=mdp.feet_contact_fixed,
+        weight=0.25,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "command_name": "base_velocity",
+            "stand_threshold": 0.1,
+            "force_threshold": 5.0,
+        },
+    )
+
+    feet_air_time = RewTerm(
+        func=mdp.feet_air_time,
+        weight=10.0,
+        params={
+            "command_name": "base_velocity",
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "threshold": 0.5,
+        },
+    )
+    
+    # stand_still = RewTerm(
+    #     func=mdp.stand_still_joint_deviation_l1,
+    #     weight=-1.0,
+    #     params={
+    #         "command_name": "base_velocity",
+    #         "command_threshold": 0.1
+    #     }
+    # )
+    
+    feet_air_time_dense = RewTerm(
+        func=mdp.feet_air_time_positive_biped,
+        weight=1.25,
+        params={
+            "command_name": "base_velocity",
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "threshold": 0.5,
+        },        
+    )
+
+    foot_stumble = RewTerm(
+        func=mdp.feet_stumble,
+        weight=-1.0,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "tangential_threshold": 2.0,
+            "normal_threshold": 1.0,
+        },
+    )
+    
+    # Penalize ankle joint limits
+    dof_pos_limits = RewTerm(
+        func=mdp.joint_pos_limits, weight=-10.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}
+    )
+    
+    energy_cost = RewTerm(
+        func=mdp.energy_cost_with_curriculum,
+        weight=-0.004,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"]),
+                "start_scale": 0.1,
+                "power": 0.8,
+                "interval_epochs": 200*24,
+                },
+    )
+
+    feet_slide = RewTerm(
+        func=mdp.feet_slide,
+        weight=-0.25,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+            "asset_cfg": SceneEntityCfg("robot", body_names=["LINK_ANKLE_ROLL_L", "LINK_ANKLE_ROLL_R"]),
+        },
+    )
+
+    dof_vel = RewTerm(
+        func=mdp.joint_vel_l2,
+        weight=-1.0e-5,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])},
+    )
+    dof_acc = RewTerm(
+        func=mdp.joint_acc_l2,
+        weight=-1.25e-8,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])},
+    )
+    action_rate = RewTerm(
+        func=mdp.action_rate_with_curriculum,
+        weight=-0.06,
+        params={"start_scale": 0.1,
+                "power": 0.8,
+                "interval_epochs": 200*24
+                },
+    )
+
+    action_smoothness = RewTerm(
+        func=mdp.action_smoothness_with_curriculum,
+        weight=-0.04,
+        params={"start_scale": 0.1,
+                "power": 0.8,
+                "interval_epochs": 200*24
+                },
+    )
+    
+    dof_torque = RewTerm(
+        func=mdp.joint_torques_l2,
+        weight=-1.0e-6,
+        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])},
+    )
+    termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0)
+
+
+
+
+
+@configclass
+class PM01Termination:
+    time_out = DoneTerm(func=mdp.time_out, time_out=True)
+    base_contact = DoneTerm(
+        func=mdp.illegal_contact,
+        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=["LINK_BASE", "LINK_KNEE_PITCH.*", ".*SHOULDER.*", ".*ELBOW.*"]), "threshold": 1.0},
+    )
+
+
+@configclass
+class ActionsCfg:
+    """Action specifications for the MDP."""
+
+    joint_pos = mdp.JointPositionActionCfg(asset_name="robot",
+                                           use_default_offset=True, 
+                                           preserve_order=True,
+                                           joint_names=PM_WAIST_DFS_JOINT_NAMES,
+                                           scale = {".*_HIP_PITCH_.*" : 0.5,
+                                                    ".*_HIP_ROLL_.*" : 0.2,
+                                                    ".*_HIP_YAW_.*" : 0.2,
+                                                    ".*_KNEE_PITCH_.*" : 0.5,
+                                                    ".*_ANKLE_PITCH_.*" : 0.5,
+                                                    ".*_ANKLE_ROLL_.*" : 0.2,
+                                                    ".*WAIST_YAW.*" : 0.2,
+                                                    ".*_SHOULDER_PITCH_.*" : 0.2,
+                                                    ".*_SHOULDER_ROLL_.*" : 0.2,
+                                                    ".*_SHOULDER_YAW_.*" : 0.2,
+                                                    ".*_ELBOW_PITCH_.*" : 0.2,
+                                                    ".*_ELBOW_YAW_.*" : 0.2
+                                                }
+                                           )
+
+
+@configclass
+class ObservationsCfg:
+    """Observation specifications for the MDP."""
+
+    @configclass
+    class PolicyCfg(ObsGroup):
+        """Observations for policy group."""
+
+        # observation terms (order preserved)
+        joint_pos = ObsTerm(
+            func=mdp.joint_pos_rel,
+            noise=Unoise(n_min=-0.01, n_max=0.01),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        joint_vel = ObsTerm(
+            func=mdp.joint_vel_rel,
+            noise=Unoise(n_min=-1.5, n_max=1.5),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        actions = ObsTerm(func=mdp.last_action,
+                          history_length=15)        
+        base_ang_vel = ObsTerm(func=mdp.base_ang_vel, 
+                               noise=Unoise(n_min=-0.2, n_max=0.2),
+                            history_length=15)
+        projected_gravity = ObsTerm(
+            func=mdp.projected_gravity,
+            noise=Unoise(n_min=-0.05, n_max=0.05),
+            history_length=15
+        )
+        velocity_commands = ObsTerm(func=mdp.generated_commands, 
+                                    params={"command_name": "base_velocity"})        
+        
+        def __post_init__(self):
+            self.enable_corruption = True
+            self.concatenate_terms = True
+    
+    @configclass
+    class CriticCfg(ObsGroup):
+
+        # observation terms (order preserved)
+        joint_pos = ObsTerm(
+            func=mdp.joint_pos_rel,
+            noise=Unoise(n_min=-0.01, n_max=0.01),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        joint_vel = ObsTerm(
+            func=mdp.joint_vel_rel,
+            noise=Unoise(n_min=-1.5, n_max=1.5),
+            params={
+                "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            },
+            history_length=15,
+        )
+        
+        actions = ObsTerm(func=mdp.last_action,
+                          history_length=15)        
+        base_ang_vel = ObsTerm(func=mdp.base_ang_vel, 
+                               noise=Unoise(n_min=-0.2, n_max=0.2),
+                            history_length=15)
+        projected_gravity = ObsTerm(
+            func=mdp.projected_gravity,
+            noise=Unoise(n_min=-0.05, n_max=0.05),
+            history_length=15
+        )
+        velocity_commands = ObsTerm(func=mdp.generated_commands, 
+                                    params={"command_name": "base_velocity"})
+
+
+        def __post_init__(self):
+            self.enable_corruption = True
+            self.concatenate_terms = True
+
+    policy: PolicyCfg = PolicyCfg()
+    critic: CriticCfg = CriticCfg()
+
+
+@configclass
+class PM01Commands:
+    """Command specifications for the MDP."""
+
+    base_velocity = mdp.UniformVelocityCommandCfg(
+        asset_name="robot",
+        resampling_time_range=(7.5, 7.5),
+        rel_standing_envs=0.1,
+        rel_heading_envs=1.0,
+        heading_command=True,
+        heading_control_stiffness=0.5,
+        debug_vis=True,
+        ranges=mdp.UniformVelocityCommandCfg.Ranges(
+            lin_vel_x=(-1.0, 1.5),
+            lin_vel_y=(-0.5, 0.5),
+            ang_vel_z=(-1.0, 1.0),
+            heading=(-3.14, 3.14),
+        ),
+    )
+
+
+VELOCITY_RANGE = {
+    "x": (-0.5, 0.5),
+    "y": (-0.5, 0.5),
+    "z": (-0.2, 0.2),
+    "roll": (-0.52, 0.52),
+    "pitch": (-0.52, 0.52),
+    "yaw": (-0.78, 0.78),
+}
+
+@configclass
+class PM01EventCfg:
+    """PM01-specific randomizations."""
+
+    # startup
+    physics_material = EventTerm(
+        func=mdp.randomize_rigid_body_material,
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
+            "static_friction_range": (0.3, 1.6),
+            "dynamic_friction_range": (0.3, 1.2),
+            "restitution_range": (0.0, 0.5),
+            "num_buckets": 64,
+        },
+    )
+
+    add_joint_default_pos = EventTerm(
+        func=mdp.randomize_joint_default_pos,
+        mode="startup",
+        params={
+            "asset_cfg": PM01_DFS_JOINT_ORDER_ASSET_CFG,
+            "pos_distribution_params": (-0.01, 0.01),
+            "operation": "add",
+        },
+    )
+
+    base_com = EventTerm(
+        func=mdp.randomize_rigid_body_com,
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names="LINK_BASE"),
+            "com_range": {"x": (-0.025, 0.025), "y": (-0.05, 0.05), "z": (-0.05, 0.05)},
+        },
+    )
+
+    # interval
+    push_robot = EventTerm(
+        func=mdp.push_by_setting_velocity,
+        mode="interval",
+        interval_range_s=(1.0, 3.0),
+        params={"velocity_range": VELOCITY_RANGE},
+    )
+
+    # rand_mass = EventTerm(
+    #     func=mdp.randomize_rigid_body_mass,
+    #     mode="startup",
+    #     params={
+    #         "asset_cfg": SceneEntityCfg("robot"),
+    #         "mass_distribution_params": (0.9, 1.1),
+    #         "operation": "scale",
+    #     },
+    # )    
+
+    reset_base = EventTerm(
+        func=mdp.reset_root_state_uniform,
+        mode="reset",
+        params={
+            "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)},
+            "velocity_range": {
+                "x": (0.0, 0.0),
+                "y": (0.0, 0.0),
+                "z": (0.0, 0.0),
+                "roll": (0.0, 0.0),
+                "pitch": (0.0, 0.0),
+                "yaw": (0.0, 0.0),
+            },
+        },
+    )
+    
+    reset_robot_joints = EventTerm(
+        func=mdp.reset_joints_by_scale,
+        mode="reset",
+        params={"position_range": (0.8, 1.2), "velocity_range": (-0.5, 0.5)},
+    )
+
+
+@configclass
+class CurriculumCfg:
+    """Curriculum terms for the MDP."""
+
+    terrain_levels = CurrTerm(func=mdp.terrain_levels_vel)
+
+
+@configclass
+class PM01FlatEnvCfg(ManagerBasedRLEnvCfg):
+    """Environment configuration directly extending the base RL env config."""
+
+    scene: PM01SceneCfg = PM01SceneCfg(num_envs=4096, env_spacing=2.5)
+    observations: ObservationsCfg = ObservationsCfg()
+    actions: ActionsCfg = ActionsCfg()
+    commands: PM01Commands = PM01Commands()
+    rewards: PM01Rewards = PM01Rewards()
+    terminations: PM01Termination = PM01Termination()
+    events: PM01EventCfg = PM01EventCfg()
+    curriculum: CurriculumCfg = CurriculumCfg()
+
+    def __post_init__(self):
+        """Apply sim wiring and curriculum toggles."""
+        # simulation settings
+        self.decimation = 5
+        self.episode_length_s = 20.0
+        self.sim.dt = 0.002
+        self.sim.render_interval = self.decimation
+        self.sim.physics_material = self.scene.terrain.physics_material
+        self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15
+        if self.scene.height_scanner is not None:
+            self.scene.height_scanner.update_period = self.decimation * self.sim.dt
+        if self.scene.contact_forces is not None:
+            self.scene.contact_forces.update_period = 0.005
+            
+
+        if getattr(self.curriculum, "terrain_levels", None) is not None:
+            if self.scene.terrain.terrain_generator is not None:
+                self.scene.terrain.terrain_generator.curriculum = True
+        else:
+            if self.scene.terrain.terrain_generator is not None:
+                self.scene.terrain.terrain_generator.curriculum = False

source/engineai_lab/tasks/velocity/mdp/__init__.py

@@ -0,0 +1,7 @@
+from isaaclab_tasks.manager_based.locomotion.velocity.mdp import *
+
+from .commands import *  # noqa: F401, F403
+from .rewards import *  # noqa: F401, F403
+from .observations import *  # noqa: F401, F403
+from .events import *  # noqa: F401, F403
+from .terminations import *  # noqa: F401, F403

source/engineai_lab/tasks/velocity/mdp/commands.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import MISSING
+
+import torch
+
+import isaaclab.utils.math as math_utils
+from isaaclab.envs.mdp.commands import UniformVelocityCommand, UniformVelocityCommandCfg
+from isaaclab.utils import configclass
+
+
+class BodyVelocityCommand(UniformVelocityCommand):
+    """Velocity command whose metrics and markers use a configured robot body."""
+
+    cfg: BodyVelocityCommandCfg
+
+    def __init__(self, cfg: BodyVelocityCommandCfg, env):
+        if not 0.0 <= cfg.rel_straight_envs <= 1.0:
+            raise ValueError(f"rel_straight_envs must be in [0, 1], got {cfg.rel_straight_envs}.")
+        super().__init__(cfg, env)
+        body_ids, body_names = self.robot.find_bodies(cfg.body_name, preserve_order=True)
+        if len(body_ids) != 1:
+            raise ValueError(
+                f"BodyVelocityCommand requires exactly one body matching {cfg.body_name!r}; "
+                f"found {body_names}."
+            )
+        self.body_id = body_ids[0]
+
+    def _resample_command(self, env_ids: Sequence[int]):
+        """Sample commands, with a configurable fraction reserved for straight walking."""
+        env_ids = torch.as_tensor(env_ids, device=self.device, dtype=torch.long)
+        super()._resample_command(env_ids)
+
+        if self.cfg.rel_straight_envs <= 0.0 or env_ids.numel() == 0:
+            return
+
+        straight_mask = torch.rand(env_ids.numel(), device=self.device) <= self.cfg.rel_straight_envs
+        straight_env_ids = env_ids[straight_mask]
+        self.vel_command_b[straight_env_ids, 1:] = 0.0
+        self.is_heading_env[straight_env_ids] = False
+
+    def _body_heading_quat_w(self) -> torch.Tensor:
+        body_quat_w = self.robot.data.body_quat_w[:, self.body_id, :]
+        _, _, yaw = math_utils.euler_xyz_from_quat(body_quat_w)
+        zeros = torch.zeros_like(yaw)
+        return math_utils.quat_from_euler_xyz(zeros, zeros, yaw - self.cfg.heading_yaw_offset)
+
+    def _body_lin_vel_heading(self) -> torch.Tensor:
+        body_lin_vel_w = self.robot.data.body_lin_vel_w[:, self.body_id, :]
+        return math_utils.quat_apply_inverse(self._body_heading_quat_w(), body_lin_vel_w)
+
+    def _update_metrics(self):
+        max_command_time = self.cfg.resampling_time_range[1]
+        max_command_step = max_command_time / self._env.step_dt
+        body_lin_vel_heading = self._body_lin_vel_heading()
+        body_yaw_rate = self.robot.data.body_ang_vel_w[:, self.body_id, 2]
+        self.metrics["error_vel_xy"] += (
+            torch.norm(self.vel_command_b[:, :2] - body_lin_vel_heading[:, :2], dim=-1) / max_command_step
+        )
+        self.metrics["error_vel_yaw"] += (
+            torch.abs(self.vel_command_b[:, 2] - body_yaw_rate) / max_command_step
+        )
+
+    def _update_command(self):
+        if self.cfg.heading_command:
+            env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
+            if len(env_ids) > 0:
+                body_quat_w = self.robot.data.body_quat_w[env_ids, self.body_id, :]
+                _, _, body_yaw = math_utils.euler_xyz_from_quat(body_quat_w)
+                body_heading = body_yaw - self.cfg.heading_yaw_offset
+                heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - body_heading)
+                self.vel_command_b[env_ids, 2] = torch.clip(
+                    self.cfg.heading_control_stiffness * heading_error,
+                    min=self.cfg.ranges.ang_vel_z[0],
+                    max=self.cfg.ranges.ang_vel_z[1],
+                )
+
+        standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
+        self.vel_command_b[standing_env_ids, :] = 0.0
+
+    def _debug_vis_callback(self, _event):
+        if not self.robot.is_initialized:
+            return
+
+        marker_pos_w = self.robot.data.body_pos_w[:, self.body_id, :].clone()
+        marker_pos_w[:, 2] += self.cfg.marker_height_offset
+        desired_pos_w = marker_pos_w.clone()
+        actual_pos_w = marker_pos_w.clone()
+        desired_pos_w[:, 2] += 0.5 * self.cfg.marker_vertical_separation
+        actual_pos_w[:, 2] -= 0.5 * self.cfg.marker_vertical_separation
+        desired_scale, desired_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
+        actual_scale, actual_quat = self._resolve_xy_velocity_to_arrow(self._body_lin_vel_heading()[:, :2])
+        self.goal_vel_visualizer.visualize(desired_pos_w, desired_quat, desired_scale)
+        self.current_vel_visualizer.visualize(actual_pos_w, actual_quat, actual_scale)
+
+    def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale
+        arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
+        arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0
+
+        heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
+        zeros = torch.zeros_like(heading_angle)
+        arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
+        arrow_quat = math_utils.quat_mul(self._body_heading_quat_w(), arrow_quat)
+        return arrow_scale, arrow_quat
+
+
+@configclass
+class BodyVelocityCommandCfg(UniformVelocityCommandCfg):
+    """Configuration for body-referenced planar velocity commands."""
+
+    class_type: type = BodyVelocityCommand
+
+    body_name: str = MISSING
+    """Body used for velocity metrics and command visualization."""
+
+    heading_yaw_offset: float = 0.0
+    """Yaw offset from the body's URDF frame to the locomotion heading frame."""
+
+    rel_straight_envs: float = 0.0
+    """Fraction of sampled environments with zero lateral and yaw commands."""
+
+    marker_height_offset: float = 0.35
+    """Vertical marker offset from the configured body origin in meters."""
+
+    marker_vertical_separation: float = 0.08
+    """Vertical separation between desired and measured velocity arrows in meters."""

source/engineai_lab/tasks/velocity/mdp/events.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import torch
+from typing import TYPE_CHECKING, Literal
+
+import isaaclab.utils.math as math_utils
+from isaaclab.assets import Articulation
+from isaaclab.envs.mdp.events import _randomize_prop_by_op,randomize_rigid_body_com
+from isaaclab.managers import SceneEntityCfg
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedEnv
+
+
+def randomize_joint_default_pos(
+    env: ManagerBasedEnv,
+    env_ids: torch.Tensor | None,
+    asset_cfg: SceneEntityCfg,
+    pos_distribution_params: tuple[float, float] | None = None,
+    operation: Literal["add", "scale", "abs"] = "abs",
+    distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
+):
+    """
+    Randomize the joint default positions which may be different from URDF due to calibration errors.
+    """
+    # extract the used quantities (to enable type-hinting)
+    asset: Articulation = env.scene[asset_cfg.name]
+
+    # save nominal value for export
+    asset.data.default_joint_pos_nominal = torch.clone(asset.data.default_joint_pos[0])
+
+    # resolve environment ids
+    if env_ids is None:
+        env_ids = torch.arange(env.scene.num_envs, device=asset.device)
+
+    # resolve joint indices
+    if asset_cfg.joint_ids == slice(None):
+        joint_ids = slice(None)  # for optimization purposes
+    else:
+        joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device)
+
+    if pos_distribution_params is not None:
+        pos = asset.data.default_joint_pos.to(asset.device).clone()
+        pos = _randomize_prop_by_op(
+            pos, pos_distribution_params, env_ids, joint_ids, operation=operation, distribution=distribution
+        )[env_ids][:, joint_ids]
+
+        if env_ids != slice(None) and joint_ids != slice(None):
+            env_ids = env_ids[:, None]
+        asset.data.default_joint_pos[env_ids, joint_ids] = pos
+        # update the offset in action since it is not updated automatically
+        env.action_manager.get_term("joint_pos")._offset[env_ids] = pos.unsqueeze(1)
+

source/engineai_lab/tasks/velocity/mdp/observations.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import torch
+from typing import TYPE_CHECKING
+
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.utils.math import euler_xyz_from_quat, quat_apply_inverse, quat_from_euler_xyz
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedEnv
+
+
+def _single_body_id(asset_cfg: SceneEntityCfg, term_name: str) -> int:
+    body_ids = asset_cfg.body_ids
+    if body_ids is None:
+        raise ValueError(f"{term_name} requires asset_cfg with exactly one body name.")
+    if isinstance(body_ids, int):
+        return body_ids
+    if isinstance(body_ids, slice) or len(body_ids) != 1:
+        raise ValueError(f"{term_name} requires exactly one body id, got {body_ids}.")
+    return body_ids[0]
+
+
+def _heading_quat_with_offset(body_quat_w: torch.Tensor, heading_yaw_offset: float = 0.0) -> torch.Tensor:
+    roll, pitch, yaw = euler_xyz_from_quat(body_quat_w)
+    zeros = torch.zeros_like(yaw)
+    return quat_from_euler_xyz(zeros, zeros, yaw - heading_yaw_offset)
+
+
+def robot_base_lin_vel_b(env: ManagerBasedEnv) -> torch.Tensor:
+    """Base linear velocity expressed in the base frame."""
+    asset = env.scene["robot"]
+    # prefer direct base-frame velocity if available
+    if getattr(asset.data, "root_lin_vel_b", None) is not None:
+        return asset.data.root_lin_vel_b
+    # fallback: rotate world velocity into base frame
+    return quat_apply_inverse(asset.data.root_quat_w, asset.data.root_lin_vel_w)
+
+
+def body_ang_vel_yaw_frame(
+    env: ManagerBasedEnv,
+    asset_cfg: SceneEntityCfg,
+    heading_yaw_offset: float = 0.0,
+) -> torch.Tensor:
+    """Body angular velocity expressed in a yaw-aligned frame for that body."""
+    body_id = _single_body_id(asset_cfg, "body_ang_vel_yaw_frame")
+    asset = env.scene[asset_cfg.name]
+    body_quat_w = asset.data.body_quat_w[:, body_id, :]
+    body_ang_vel_w = asset.data.body_ang_vel_w[:, body_id, :]
+    return quat_apply_inverse(_heading_quat_with_offset(body_quat_w, heading_yaw_offset), body_ang_vel_w)
+
+
+def body_projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+    """Gravity projection in a configured body's local frame."""
+    body_id = _single_body_id(asset_cfg, "body_projected_gravity")
+    asset = env.scene[asset_cfg.name]
+    return quat_apply_inverse(asset.data.body_quat_w[:, body_id, :], asset.data.GRAVITY_VEC_W)

source/engineai_lab/tasks/velocity/mdp/terminations.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+import torch
+from isaaclab.envs import mdp
+from isaaclab.managers import SceneEntityCfg
+
+
+def illegal_contact(env, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+    """Terminate on illegal contacts and return the mask on the RL env device."""
+    return mdp.illegal_contact(env, threshold=threshold, sensor_cfg=sensor_cfg).to(env.device)

source/engineai_lab/utils/AMP_data_loader.py

@@ -0,0 +1,169 @@
+import math
+import numpy as np
+import os
+import torch
+import yaml
+from collections.abc import Iterator
+from pathlib import Path
+
+from isaaclab.utils.math import subtract_frame_transforms, quat_apply, quat_inv
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+
+
+def _resolve_repo_path(path_like: str) -> Path:
+    """Resolve path relative to the repository root when not absolute."""
+    path = Path(path_like).expanduser()
+    if not path.is_absolute():
+        path = (_REPO_ROOT / path).resolve()
+    return path
+
+
+def _load_motion_list_from_yaml(yaml_path: Path) -> list[str]:
+    """Parse YAML motion config into absolute file paths."""
+    if not yaml_path.is_file():
+        raise AssertionError(f"Invalid YAML path: {yaml_path}")
+
+    with yaml_path.open("r", encoding="utf-8") as yaml_file:
+        data = yaml.safe_load(yaml_file) or {}
+    motions = data.get("motions", [])
+    base_dir = yaml_path.parent
+    paths: list[str] = []
+
+    for entry in motions:
+        if "file" in entry:
+            file_path = (base_dir / entry["file"]).resolve()
+            if not file_path.is_file():
+                raise AssertionError(f"Motion file not found: {file_path}")
+            paths.append(str(file_path))
+        elif "folder" in entry:
+            folder_path = (base_dir / entry["folder"]).resolve()
+            if not folder_path.is_dir():
+                raise AssertionError(f"Motion folder not found: {folder_path}")
+            npz_files = sorted(p for p in folder_path.iterdir() if p.suffix == ".npz")
+            if not npz_files:
+                raise AssertionError(f"No .npz files found in folder: {folder_path}")
+            paths.extend(str(file_path) for file_path in npz_files)
+    return paths
+
+class AMPDataLoader:
+    def __init__(
+        self,
+        motion_file: str | list[str],
+        device: str = "cpu",
+        history_length: int = 5,
+    ):
+        assert history_length >= 1, "history_length must be positive"
+
+        file_list: list[str] = []
+        if isinstance(motion_file, str):
+            motion_path = _resolve_repo_path(motion_file)
+            if motion_path.suffix == ".yaml":
+                file_list.extend(_load_motion_list_from_yaml(motion_path))
+            elif motion_path.is_file() and motion_path.suffix == ".npz":
+                file_list.append(str(motion_path))
+            elif motion_path.is_dir():
+                for file_name in os.listdir(motion_path):
+                    full_path = motion_path / file_name
+                    if full_path.suffix == ".npz" and full_path.is_file():
+                        file_list.append(str(full_path))
+            else:
+                raise AssertionError(f"Invalid motion source: {motion_file}")
+        elif isinstance(motion_file, list):
+            for file_name in motion_file:
+                motion_path = _resolve_repo_path(file_name)
+                if motion_path.suffix == ".yaml":
+                    file_list.extend(_load_motion_list_from_yaml(motion_path))
+                elif motion_path.is_file() and motion_path.suffix == ".npz":
+                    file_list.append(str(motion_path))
+
+        assert len(file_list) > 0, f"No valid motion data found in: {motion_file}"
+        print("\n=========== AMP Motion File List ===========")
+        for idx, path in enumerate(file_list):
+            print(f"{idx + 1:2d}. {path}")
+        print(f"=========== Total: {len(file_list)} files ===========\n")
+
+        fps_list: list[float] = []
+        joint_pos_list: list[torch.Tensor] = []
+        joint_vel_list: list[torch.Tensor] = []
+        body_pos_w_list: list[torch.Tensor] = []
+        body_quat_w_list: list[torch.Tensor] = []
+        body_lin_vel_w_list: list[torch.Tensor] = []
+        body_ang_vel_w_list: list[torch.Tensor] = []
+
+        for file_path in file_list:
+            try:
+                data = np.load(file_path)
+                fps_list.append(float(data["fps"]))
+                joint_pos_list.append(torch.tensor(data["joint_pos"], dtype=torch.float32, device=device))
+                joint_vel_list.append(torch.tensor(data["joint_vel"], dtype=torch.float32, device=device))
+                body_pos_w_list.append(torch.tensor(data["body_pos_w"], dtype=torch.float32, device=device))
+                body_quat_w_list.append(torch.tensor(data["body_quat_w"], dtype=torch.float32, device=device))
+                body_lin_vel_w_list.append(torch.tensor(data["body_lin_vel_w"], dtype=torch.float32, device=device))
+                body_ang_vel_w_list.append(torch.tensor(data["body_ang_vel_w"], dtype=torch.float32, device=device))
+            except Exception as exc:  # noqa: BLE001
+                print(f"Warning: Could not load {file_path}: {exc}")
+
+        assert len(joint_pos_list) > 0, "Failed to load any motion data"
+        self.fps = torch.tensor(fps_list, dtype=torch.float32, device=device)
+        self.joint_pos = torch.cat(joint_pos_list, dim=0)
+        self.joint_vel = torch.cat(joint_vel_list, dim=0)
+        self.body_pos_w = torch.cat(body_pos_w_list, dim=0)
+        self.body_quat_w = torch.cat(body_quat_w_list, dim=0)
+        self.body_lin_vel_w = torch.cat(body_lin_vel_w_list, dim=0)
+        self.body_ang_vel_w = torch.cat(body_ang_vel_w_list, dim=0)
+        self.time_step_total = self.joint_pos.shape[0]
+        self.body_pos_b = torch.zeros_like(self.body_pos_w, device=device)
+        self.body_quat_b = torch.zeros_like(self.body_quat_w, device=device)
+        self.body_lin_vel_b = torch.zeros_like(self.body_lin_vel_w, device=device)
+        self.body_ang_vel_b = torch.zeros_like(self.body_ang_vel_w, device=device)
+        self.projected_gravity_b = torch.zeros((self.time_step_total, 3), dtype=torch.float32, device=device)
+        self.num_bodies = self.body_pos_w.shape[1]
+        self.history_length = history_length
+
+        for i in range(self.time_step_total):            
+            body_pos_w_t = self.body_pos_w[i].unsqueeze(0)  # (1, B, 3)
+            body_quat_w_t = self.body_quat_w[i].unsqueeze(0)  # (1, B, 4)
+            achor_pos_w_t = body_pos_w_t[:, 0:1, :]  # (1, 1, 3) 
+            achor_quat_w_t = body_quat_w_t[:, 0:1, :]  # (1, 1, 4)
+            pos_b_t, quat_b_t = subtract_frame_transforms(
+                achor_pos_w_t.repeat(1, self.num_bodies, 1),
+                achor_quat_w_t.repeat(1, self.num_bodies, 1),
+                body_pos_w_t,
+                body_quat_w_t,
+            )
+            
+            self.body_pos_b[i] = pos_b_t.squeeze(0)
+            self.body_quat_b[i] = quat_b_t.squeeze(0)
+            # rotate velocities into anchor frame with inverse; expand anchor quat per body for broadcasting
+            anchor_quat_broadcast = quat_inv(achor_quat_w_t.repeat(1, self.num_bodies, 1)).reshape(-1, 4)
+            body_lin_vel = self.body_lin_vel_w[i].unsqueeze(0).reshape(-1, 3)
+            self.body_lin_vel_b[i] = quat_apply(anchor_quat_broadcast, body_lin_vel).reshape(self.num_bodies, 3)
+        self.base_lin_vel_b = self.body_lin_vel_b[:, 0, :]
+
+    # Ugly mini batch generator; the observation acquirement need to be re-designed
+    def mini_batch_generator(self, num_mini_batches, num_epoches) -> Iterator[torch.Tensor]:
+        """Generate mini-batches of motion data."""
+        num_samples = self.joint_pos.shape[0]
+        batch_size = math.ceil(num_samples / num_mini_batches)
+        indices = torch.randperm(num_samples, device=self.joint_pos.device)
+        # history is ordered old -> new; offsets are negative to zero so the last frame is "current"
+        history_offsets = torch.arange(-(self.history_length - 1), 1, device=self.joint_pos.device)
+
+        def gather_frame_features(idxs: torch.Tensor) -> list[torch.Tensor]:
+            """Collect per-frame features for the given indices."""
+            pos = self.joint_pos[idxs] * 9
+            base_lin = self.base_lin_vel_b[idxs] * 7
+            return [pos, base_lin]
+
+        for _ in range(num_epoches):
+            for i in range(0, num_samples, batch_size):
+                batch_indices = indices[i:i + batch_size]
+
+                frame_features = []
+                # history direction: old -> new (last frame aligns with batch_indices)
+                for offset in history_offsets:
+                    idxs = torch.clamp(batch_indices + offset, min=0, max=num_samples - 1)
+                    frame_features.extend(gather_frame_features(idxs))
+
+                yield torch.cat(frame_features, dim=-1)

source/engineai_lab/utils/AMP_discriminator.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from dataclasses import dataclass, field
+from rsl_rl.modules import MLP, EmpiricalNormalization
+from rsl_rl.utils import resolve_nn_activation
+
+class Discriminator(nn.Module):
+    """
+    Simple feedforward neural network as a discriminator for Adversarial Motion Prior.
+    """
+
+    def __init__(
+        self,
+        input_dim_per_frame: int = 58,
+        input_history_length: int = 1,
+        hidden_dims: list[int] = [256, 128],
+        activation: str = "relu",
+        feature_normalization: bool = False,  # if True, normalize input features with EmpiricalNormalization
+        device: str = "cpu"
+    ):
+        super().__init__()
+
+        self.device = device
+        curr_in_dim = input_dim_per_frame * input_history_length
+        print("Discriminator input dim:", curr_in_dim)
+        self.frame_size = input_dim_per_frame
+        self.history_length = input_history_length
+        
+        self.activation = resolve_nn_activation(activation) #type: ignore
+        
+        layers = []
+        for hidden_dim in hidden_dims:
+            layers.append(nn.Linear(curr_in_dim, hidden_dim))
+            layers.append(self.activation) # use resolved activation
+            curr_in_dim = hidden_dim
+        self.model = nn.Sequential(*layers).to(self.device) # type: ignore
+        self.linear_layer = nn.Linear(hidden_dims[-1], 1).to(self.device)
+
+        self.feature_normalization = feature_normalization
+        if self.feature_normalization:
+            self.feature_norm = EmpiricalNormalization(shape=(self.frame_size,)).to(self.device)
+
+    def normalize_input(self, x):
+        if self.feature_normalization:
+            # avoid in-place on a leaf tensor by normalizing frame slices and re-concatenating
+            frames = torch.split(x, self.frame_size, dim=1)
+            norm_frames = [self.feature_norm(frame) for frame in frames]
+            x = torch.cat(norm_frames, dim=1)
+        return x
+
+    def forward(self, x):
+        assert self.history_length * self.frame_size == x.shape[1], \
+            f"Input feature dimension {x.shape[1]} does not match expected size {self.history_length * self.frame_size}"
+        if self.feature_normalization:
+            x = self.normalize_input(x)
+        return self.linear_layer(self.model(x)).squeeze(-1)
+    
+    # TODO: normalize feature on positive or negative samples?
+    def update_normalization(self, x):
+        if self.feature_normalization:
+            for i in range(self.history_length):
+                start_idx = i * self.frame_size
+                end_idx = (i + 1) * self.frame_size
+                self.feature_norm.update(x[:, start_idx:end_idx])
+
+    @torch.no_grad()
+    def get_amp_reward(self, x):
+        self.eval()
+        d = self.forward(x)
+        r = torch.clamp(1 - 0.25 * (d - 1) ** 2, min=0)
+        self.train()
+        return r
+
+    def compute_grad_pen(self, expert_data, lambda_=10):
+        expert_data.requires_grad = True
+        disc = self.forward(expert_data)
+        ones = torch.ones(disc.size(), device=disc.device)
+        grad = torch.autograd.grad(
+            outputs=disc, inputs=expert_data,
+            grad_outputs=ones, create_graph=True,
+            retain_graph=True, only_inputs=True)[0]
+
+        # Enforce that the grad norm approaches 0.
+        grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean()
+        return grad_pen
+