robomotic/mjlab-policies

MJLab Policies

This repository contains trained reinforcement learning policies for robotic tasks using MJLab and RSL-RL.

Repository Structure

Policies are organized by task and training run:

<task-name>/
└── <timestamp>/
    ├── model_final.pt       # PyTorch checkpoint (final iteration)
    ├── policy.onnx           # Exported ONNX policy
    ├── agent.yaml           # Agent configuration
    ├── env.yaml             # Environment configuration
    ├── mjlab.diff           # Git diff for reproducibility
    └── README.md            # Run-specific metadata and details

Downloading a Policy

Using Git LFS

# Clone the repository
git lfs install
git clone https://huggingface.co/robomotic/mjlab-policies
cd mjlab-policies

# Navigate to the desired task and run
cd <task-name>/<timestamp>/

Using Hugging Face CLI

# Download a specific checkpoint
huggingface-cli download robomotic/mjlab-policies <task-name>/<timestamp>/model_final.pt

# Download the ONNX model
huggingface-cli download robomotic/mjlab-policies <task-name>/<timestamp>/policy.onnx

Using Python

from huggingface_hub import hf_hub_download

checkpoint_path = hf_hub_download(
    repo_id="robomotic/mjlab-policies",
    filename="<task-name>/<timestamp>/model_final.pt"
)

Loading a Policy

With MjLab

# Install mjlab
git clone https://github.com/mujocolab/mjlab.git
cd mjlab
uv sync

# Play the policy
uv run play --task <task-name> --checkpoint path/to/model_final.pt

PyTorch Checkpoint

import torch

checkpoint = torch.load("model_final.pt", map_location="cpu")
actor_state_dict = checkpoint["actor_state_dict"]
# Load into your MJLab runner

ONNX Policy

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("policy.onnx")
obs = np.zeros((1, observation_dim), dtype=np.float32)
action = session.run(None, {session.get_inputs()[0].name: obs})[0]

Available Tasks

Task	Runs	Description
See directory listing above for available tasks and timestamps

Citation

If you use these policies in your research, please cite MJLab:

@software{mjlab2024,
  author = {The MjLab Developers},
  title = {MjLab: Isaac Lab API, powered by MuJoCo-Warp},
  url = {https://github.com/mujocolab/mjlab},
  year = {2024}
}

Questions?

For questions about MJLab or these policies, please open an issue on the MJLab GitHub repository.