Add mjlab exporter (#47)

# Pull Request

### What change is being made

Please provide a detailed description of WHAT change is being made such
that the reader is able to understand the change holistically.

### Why this change is being made

Please provide the rationale behind the change and additional context like links to documents, related work or tickets.

### Tested

Please provide a description how this change was tested, e.g. unit tests, hardware tests or commands you run.

GitOrigin-RevId: a4174dfb9929d45864f0090b23fdf9fe8e9c72bb

Author: jbruedigam-bdai

.github/workflows/test.yml

@@ -90,7 +90,7 @@ jobs:
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: isaaclab
 
-      - name: Run Isaac Lab Export Test
+      - name: Run IsaacLab Export Test
         env:
           OMNI_KIT_ACCEPT_EULA: yes
         run: pixi run -e isaaclab export-isaaclab-ci
@@ -116,3 +116,41 @@ jobs:
         env:
           OMNI_KIT_ACCEPT_EULA: yes
         run: pixi run -e isaaclab test
+
+  test-mjlab-export:
+    name: Test Export mjlab
+    runs-on: github-gpu-runner
+    timeout-minutes: 60
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+
+      - name: Setup pixi
+        uses: prefix-dev/setup-pixi@v0.9.4
+        with:
+          pixi-version: latest
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
+          environments: mjlab
+
+      - name: Run mjlab Export Test
+        run: pixi run -e mjlab export-mjlab-ci
+
+  test-frameworks-mjlab:
+    name: Test Frameworks - mjlab
+    runs-on: github-gpu-runner
+    timeout-minutes: 60
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+
+      - name: Setup pixi
+        uses: prefix-dev/setup-pixi@v0.9.4
+        with:
+          pixi-version: latest
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
+          environments: mjlab
+
+      - name: Run Frameworks mjlab Tests
+        run: pixi run -e mjlab test

.ruff.toml

@@ -19,14 +19,8 @@ select = [
 ignore = [
     "E501",  # line-too-long (formatter handles this)
     "E402",  # module-level-import-not-at-top-of-file (needed for some frameworks)
-    "F401",  # unused-import (may be for re-export)
-    "F841",  # unused-variable (sometimes intentional)
 ]
 
-[lint.per-file-ignores]
-"__init__.py" = ["F401"]  # Allow unused imports in __init__.py
-"examples/exporter_scripts/export_isaaclab.py" = ["E402", "F401", "F403"]  # Allow import flexibility for isaaclab
-
 [format]
 quote-style = "double"
 indent-style = "space"

.vscode/launch.json

@@ -8,11 +8,20 @@
             "name": "[Examples] Export IsaacLab task.",
             "type": "debugpy",
             "request": "launch",
-            "program": "${workspaceFolder}/examples/exporter_scripts/export_isaaclab.py",
+            "program": "${workspaceFolder}/examples/exporter_scripts/isaaclab/export.py",
             "console": "integratedTerminal",
             "python": "${workspaceFolder}/.pixi/envs/isaaclab/bin/python",
             "justMyCode": false
         },
+        {
+            "name": "[Examples] Export mjlab task.",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/examples/exporter_scripts/mjlab/export.py",
+            "console": "integratedTerminal",
+            "python": "${workspaceFolder}/.pixi/envs/mjlab/bin/python",
+            "justMyCode": false,
+        },
         {
             "name": "[Core] Tests.",
             "type": "debugpy",
@@ -27,6 +36,6 @@
             "env": {
                 "PYTHONPATH": "${workspaceFolder}"
             }
-        },
+        }
     ]
 }

control/matcher.cpp

@@ -295,9 +295,13 @@ std::vector<std::unique_ptr<Input>> DepthImageMatcher::createInputs() const {
 // ---------------  Body matchers ------------------------------
 bool BodyPositionMatcher::matches(const Match& maybe_match) {
   std::smatch match;
-  std::regex pattern = std::regex(
-      fmt::format("obj\\.({})\\.bodies\\.({})\\.pos_b_rt_w_in_w", alphanumeric, alphanumeric));
+  std::regex pattern =
+      std::regex(fmt::format("obj\\.({})\\.({})\\.pos_b_rt_w_in_w", alphanumeric, alphanumeric));
   if (std::regex_match(maybe_match.name, match, pattern) && match.size() > 2) {
+    // Exclude "base" - that's handled by BasePositionMatcher
+    if (match[2].str() == "base") {
+      return false;
+    }
     found_matches_[match[2].str()] = maybe_match;
     return true;
   }
@@ -315,8 +319,12 @@ std::vector<std::unique_ptr<Input>> BodyPositionMatcher::createInputs() const {
 bool BodyOrientationMatcher::matches(const Match& maybe_match) {
   std::smatch match;
   std::regex pattern =
-      std::regex(fmt::format("obj\\.({})\\.bodies\\.({})\\.w_Q_b", alphanumeric, alphanumeric));
+      std::regex(fmt::format("obj\\.({})\\.({})\\.w_Q_b", alphanumeric, alphanumeric));
   if (std::regex_match(maybe_match.name, match, pattern) && match.size() > 2) {
+    // Exclude "base" - that's handled by BaseOrientationMatcher
+    if (match[2].str() == "base") {
+      return false;
+    }
     found_matches_[match[2].str()] = maybe_match;
     return true;
   }

control/test/components_test.cpp

@@ -133,7 +133,7 @@ TEST_F(OnnxComponentsTest, JointVelocityInput_InitAndRead) {
 }
 
 TEST_F(OnnxComponentsTest, BodyOrientationInput_InitAndRead) {
-  BodyOrientationInput body_input("obj.box1.bodies.box.w_Q_b", "test_body");
+  BodyOrientationInput body_input("obj.box1.box.w_Q_b", "test_body");
 
   // Test initialization
   EXPECT_CALL(state_mock_, initBodyOrientationW("test_body")).WillOnce(Return(true));

control/test/testdata/test_onnx_generator.py

@@ -40,8 +40,8 @@
     "sensor.ray_caster.trail.b",
     "sensor.depth_image.one",
     # body
-    "obj.box1.bodies.box.pos_b_rt_w_in_w",
-    "obj.box1.bodies.box.w_Q_b",
+    "obj.box1.box.pos_b_rt_w_in_w",
+    "obj.box1.box.w_Q_b",
     # memory
     "memory.output.joint_targets.jt1.pos.in",
     # step count

docs/api/frameworks/isaaclab.rst

@@ -1,10 +1,10 @@
-Isaac Lab Framework
+IsaacLab Framework
 ===================
 
-Isaac Lab integration for RL policy export.
+IsaacLab integration for RL policy export.
 
 This module provides the ``IsaacLabExportableEnvironment`` class and utilities
-for exporting policies trained in NVIDIA Isaac Lab to ONNX format.
+for exporting policies trained in NVIDIA IsaacLab to ONNX format.
 
 .. automodule:: exploy.exporter.frameworks.isaaclab
    :members:

…rter_scripts/isaaclab/export_isaaclab.py …ples/exporter_scripts/isaaclab/export.pyexamples/exporter_scripts/isaaclab/export_isaaclab.py renamed to examples/exporter_scripts/isaaclab/export.py examples/exporter_scripts/isaaclab/export_isaaclab.py renamed to examples/exporter_scripts/isaaclab/export.py

@@ -17,14 +17,14 @@
 
 def make_simulation_app() -> tuple[SimulationApp, argparse.Namespace]:
     # Create argument parser for headless mode
-    parser = argparse.ArgumentParser(description="Export Isaac Lab environment to ONNX")
+    parser = argparse.ArgumentParser(description="Export IsaacLab environment to ONNX")
 
     # Add custom arguments
     parser.add_argument(
         "--task",
         type=str,
-        default="Isaac-Velocity-Rough-G1-Play-v0",
-        help="Name of the Isaac Lab task to export (default: Isaac-Velocity-Rough-G1-Play-v0)",
+        default="IsaacLab-Velocity-Rough-G1-Play-v0",
+        help="Name of the IsaacLab task to export (default: IsaacLab-Velocity-Rough-G1-Play-v0)",
     )
     parser.add_argument(
         "--pause-on-failure",
@@ -61,20 +61,14 @@ def make_simulation_app() -> tuple[SimulationApp, argparse.Namespace]:
 from exploy.exporter.core.evaluator import evaluate
 from exploy.exporter.core.exporter import export_environment_as_onnx
 from exploy.exporter.core.session_wrapper import SessionWrapper
-from exploy.exporter.frameworks.isaaclab import (
-    environments,  # noqa: F401
-    inputs,
-    memory,
-    outputs,
-)
-from exploy.exporter.frameworks.isaaclab.actor import make_exportable_actor
+from exploy.exporter.frameworks.isaaclab import environments  # noqa: F401
 from exploy.exporter.frameworks.isaaclab.env import IsaacLabExportableEnvironment
+from exploy.exporter.frameworks.manager_based import inputs, memory, outputs
+from exploy.exporter.frameworks.manager_based.actor import make_exportable_actor
 
 
-def export_isaaclab(
-    task_name: str = "Isaac-Velocity-Rough-G1-Play-v0", pause_on_failure: bool = False
-):
-    """Test Isaac Lab ONNX export and evaluation pipeline."""
+def export(task_name: str = "Isaac-Velocity-Rough-G1-Play-v0", pause_on_failure: bool = False):
+    """Test IsaacLab ONNX export and evaluation pipeline."""
     test_dir = pathlib.Path(__file__).parent / "exporter_tests"
 
     task_device = "cpu"
@@ -93,50 +87,42 @@ def export_isaaclab(
     onnx_export_dir = test_dir
     onnx_export_file = "test_export.onnx"
 
-    exportable_env = IsaacLabExportableEnvironment(env.unwrapped)
+    unwrapped_env = env.unwrapped
+    exportable_env = IsaacLabExportableEnvironment(unwrapped_env)
 
     # Get the policy and its normalizer.
     alg: PPO = runner.alg
     assert isinstance(alg, PPO), f"Expected PPO algorithm, got: {type(alg).__name__}"
     actor = make_exportable_actor(exportable_env, alg.policy, device=task_device)
 
-    articulations = env.unwrapped.scene.articulations
+    articulations = unwrapped_env.unwrapped.scene.articulations
     context_manager = exportable_env.context_manager()
 
-    inputs.add_base_vel(
-        articulations=articulations,
-        context_manager=context_manager,
-    )
-
-    inputs.add_body_pos_and_quat(
-        articulations=articulations,
-        context_manager=context_manager,
-    )
+    inputs.add_base_vel(articulations, context_manager)
 
-    inputs.add_commands(
-        command_manager=env.unwrapped.command_manager,
-        context_manager=context_manager,
-    )
+    inputs.add_body_pos_and_quat(articulations, context_manager)
 
-    inputs.add_joint_pos_and_vel(
-        articulations=articulations,
-        context_manager=context_manager,
+    inputs.add_command(
+        unwrapped_env,
+        context_manager,
+        command_name="base_velocity",
+        command_type="se2_velocity",
     )
 
-    inputs.add_sensor_inputs(
-        sensors=env.unwrapped.scene.sensors,
-        context_manager=context_manager,
-    )
+    inputs.add_joint_pos_and_vel(articulations, context_manager)
 
-    memory.add_memory(
-        env=env.unwrapped,
-        context_manager=context_manager,
-    )
+    for sensor_name, sensor in unwrapped_env.scene.sensors.items():
+        inputs.add_sensor_input(sensor_name, sensor, context_manager)
 
-    outputs.add_outputs(
-        action_manager=env.unwrapped.action_manager,
-        context_manager=context_manager,
-    )
+    memory.add_memory(unwrapped_env, context_manager, attr_name="action")
+    for action_term_name in unwrapped_env.action_manager.active_terms:
+        memory.add_memory(
+            unwrapped_env,
+            context_manager,
+            attr_name="processed_actions",
+            action_term_name=action_term_name,
+        )
+        outputs.add_output(unwrapped_env, context_manager, action_term_name=action_term_name)
 
     export_environment_as_onnx(
         env=exportable_env,
@@ -181,7 +167,7 @@ def export_isaaclab(
     import sys
 
     try:
-        export_isaaclab(task_name=args.task, pause_on_failure=args.pause_on_failure)
+        export(task_name=args.task, pause_on_failure=args.pause_on_failure)
     except Exception as e:
         print(f"❌ Test ERROR: {e}")
         import traceback