Add support for mjlab

### What change is being made

Add support for mjlab similar to the IsaacLab implementation.

### Why this change is being made

Make exploy more useful.

### Tested

Extensive test coverage.

Co-authored-by: jbruedigam-bdai <jbruedigam@rai-inst.com>

---------

Co-authored-by: jbruedigam-bdai <jbruedigam@rai-inst.com>
Co-authored-by: Jan Bruedigam <jbruedigam@theaiinstitute.com>
Co-authored-by: Dario Bellicoso <dbellicoso@theaiinstitute.com>
GitOrigin-RevId: 80b0b81d04ce76362c3c830148c390509eea9920

Author: 3 people

.vscode/launch.json

@@ -13,6 +13,15 @@
       "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_mjlab.py",
+      "console": "integratedTerminal",
+      "python": "${workspaceFolder}/.pixi/envs/mjlab/bin/python",
+      "justMyCode": false
+    },
     {
       "name": "[Core] Tests.",
       "type": "debugpy",

README.md

@@ -36,9 +36,10 @@ and maximizes confidence that your policy will behave identically in simulation
 - **C++ Controller with ONNX Runtime Integration**:
   Deploy trained policies using ONNX Runtime for real-time policy execution.
 - **Multi-Framework Support**:
-  Built-in support for [IsaacLab][isaaclab], with an architecture designed for extensible framework integration.
+  Built-in support for [IsaacLab][isaaclab] and [MjLab][mjlab], with an architecture designed for extensible framework integration.
 
 [isaaclab]: https://github.com/isaac-sim/IsaacLab
+[mjlab]: https://github.com/mujocolab/mjlab
 
 ## Project Structure
 
@@ -84,6 +85,12 @@ To install with additional dependencies for `IsaacLab` integration:
 pip install "exploy[isaaclab]@git+https://github.com/rai-opensource/exploy.git"
 ```
 
+To install with additional dependencies for `MjLab` integration:
+
+```bash
+pip install "exploy[mjlab]@git+https://github.com/rai-opensource/exploy.git"
+```
+
 ### C++ Controller Library
 
 To consume the C++ controller library, use plain CMake commands to build and install it on your system (or to a custom prefix):
@@ -130,16 +137,30 @@ This ensures the Exploy C++ controller library is built and exposed correctly as
 > [!NOTE]
 > This example assumes you have an NVIDIA GPU available.
 
-You can export an ONNX policy for an existing [IsaacLab][isaaclab] tasks:
+You can export an ONNX policy for an existing [IsaacLab][isaaclab] task:
 
 ```bash
 pixi run export-isaaclab
 ```
 
+Or for an existing [MjLab][mjlab] task:
+
+```bash
+pixi run export-mjlab
+```
+
 Then, you can run the [C++ example][cpp_example] that loads and executes the exported policy:
 
+For IsaacLab:
+
 ```bash
-pixi run run-cpp-example
+pixi run run-cpp-example-isaaclab
+```
+
+For MjLab:
+
+```bash
+pixi run run-cpp-example-mjlab
 ```
 
 ## Development
@@ -179,11 +200,11 @@ The editor will automatically detect the Pixi environment.
 - By default, the Python environment is configured for the `exploy.exporter.core` subpackage.
 
 If you are working on a machine with an NVIDIA GPU, you can switch to the environment
-for the `exploy.exporter.frameworks.isaaclab` component:
+for the `exploy.exporter.frameworks.isaaclab` or `exploy.exporter.frameworks.mjlab` component:
 
 - Click the Python icon in the left sidebar.
 - In the `Environment Managers` panel, locate `Pixi`.
-- Select the `exploy/isaaclab` entry and set it as the project environment.
+- Select the `exploy/isaaclab` or `exploy/mjlab` entry and set it as the project environment.
 
 #### 4. Available Tasks
 
@@ -240,6 +261,9 @@ pip install -e .
 
 # Install the isaaclab component
 pip install -e ".[isaaclab]"
+
+# Install the mjlab component
+pip install -e ".[mjlab]"
 ```
 
 Using this method, you are responsible for ensuring the consistency of the dependencies in your environment,

docs/api/frameworks/mjlab.rst

@@ -0,0 +1,93 @@
+MjLab Framework
+===============
+
+This section contains documentation for the MjLab framework adapter, which provides tools for
+exporting policies from MjLab to ONNX format using exploy.
+
+
+Environment
+-----------
+
+The exportable environment adapter for MjLab.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.env
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Actor
+-----
+
+Actor wrappers for RSL-RL v5 policy networks.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.actor
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Inputs
+------
+
+Utilities for registering sensor and command inputs from MjLab managers.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.inputs
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Outputs
+-------
+
+Utilities for registering action outputs from MjLab managers.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.outputs
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Memory
+------
+
+Utilities for registering memory components from MjLab managers.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.memory
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Data Sources
+------------
+
+Adaptor classes that mirror MjLab data interfaces while managing their own
+tensor data for ONNX export.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.entity_data
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: exploy.exporter.frameworks.mjlab.raycaster_data
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: exploy.exporter.frameworks.mjlab.sensor_proxy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Utilities
+---------
+
+Helper functions for MjLab integration.
+
+.. automodule:: exploy.exporter.frameworks.mjlab.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/index.rst

@@ -22,6 +22,7 @@ Framework-specific implementations.
    :maxdepth: 2
 
    frameworks/isaaclab
+   frameworks/mjlab
 
 C++ API
 -------

docs/conf.py

@@ -95,6 +95,8 @@
     "isaaclab_rl",
     "isaaclab_tasks",
     "isaacsim",
+    "mjlab",
+    "tensordict",
     "rsl_rl",
 ]