Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
2026-08-28 15:42:17 +08:00 · 2026-08-28 15:42:17 +08:00 · 32a241c28f
commit 32a241c28f
513 changed files with 120771 additions and 0 deletions
--- a/.claude/commands/commit-push-pr.md
+++ b/.claude/commands/commit-push-pr.md
@ -0,0 +1,19 @@
+---
+allowed-tools: Bash(git checkout --branch:*), Bash(git add:*), Bash(git status:*), Bash(git push:*), Bash(git commit:*), Bash(gh pr create:*)
+description: Commit, push, and open a PR
+---
+
+## Context
+
+- Current git status: !`git status`
+- Current git diff (staged and unstaged changes): !`git diff HEAD`
+- Current branch: !`git branch --show-current`
+
+## Your task
+
+Based on the above changes:
+1. Create a new branch if on main
+2. Create a single commit with an appropriate message
+3. Push the branch to origin
+4. Create a pull request using `gh pr create`
+5. You have the capability to call multiple tools in a single response. You MUST do all of the above in a single message. Do not use any other tools or do anything else. Do not send any other text or messages besides these tool calls.
--- a/.claude/commands/update-mjwarp.md
+++ b/.claude/commands/update-mjwarp.md
@ -0,0 +1,18 @@
+---
+allowed-tools: Bash(uv lock), Bash(git checkout:*), Bash(git add:*), Bash(git status:*), Bash(git push:*), Bash(git commit:*), Bash(gh pr create:*), Edit, Read
+description: Update the mujoco-warp dependency to a given commit
+---
+
+Update the mujoco-warp dependency to commit $ARGUMENTS.
+
+Steps:
+1. Read `pyproject.toml` and find the `mujoco-warp` line under `[tool.uv.sources]`.
+2. Use Edit to replace the current `rev = "..."` value with `rev = "$ARGUMENTS"` on that line.
+3. Run `uv lock` to regenerate the lockfile.
+4. Create and switch to a new branch named `update-mjwarp/<first-8-chars-of-hash>` (e.g. `update-mjwarp/e28c6038`).
+5. Stage `pyproject.toml` and `uv.lock`, then commit with message: `Update mujoco-warp to <first-8-chars-of-hash>`.
+6. Push the branch and open a PR with title `Update mujoco-warp to <first-8-chars-of-hash>`.
+
+Important:
+- The commit hash is required. If `$ARGUMENTS` is empty, ask the user for a commit hash.
+- Do NOT modify anything else in `pyproject.toml`.
--- a/.claude/settings.json
+++ b/.claude/settings.json
@ -0,0 +1,33 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(make:*)",
+      "Bash(uv run:*)",
+      "Bash(uv lock:*)",
+      "Bash(uv sync:*)",
+      "Bash(uv add:*)",
+      "Bash(git:*)",
+      "Bash(gh:*)",
+      "WebSearch",
+      "Skill(commit-push-pr)",
+      "Skill(pr-review-toolkit:review-pr)"
+    ]
+  },
+  "hooks": {
+    "PostToolUse": [
+      {
+        "matcher": "Write|Edit",
+        "hooks": [
+          {
+            "type": "command",
+            "command": "uv run ruff format"
+          }
+        ]
+      }
+    ]
+  },
+  "enabledPlugins": {
+    "code-simplifier@claude-plugins-official": true,
+    "pr-review-toolkit@claude-plugins-official": true
+  }
+}
--- a/.dockerignore
+++ b/.dockerignore
@ -0,0 +1,34 @@
+# Large runtime directories
+.venv/
+logs/
+wandb/
+artifacts/
+benchmark_results/
+dist/
+
+# Build/cache
+__pycache__/
+*.pyc
+.ruff_cache/
+.pytest_cache/
+.uv-cache/
+*.egg-info/
+
+# Git/CI
+.git/
+.github/
+.gitignore
+.pre-commit-config.yaml
+
+# IDE/local
+.vscode/
+.claude/
+notebooks/
+
+# Docker
+Dockerfile
+.dockerignore
+
+# Docs build artifacts
+docs/source/_build/
+docs/source/generated/
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@ -0,0 +1,156 @@
+name: tests
+
+on:
+  push:
+    branches: [main]
+    paths-ignore:
+      - '**.md'
+      - '**.rst'
+      - 'docs/**'
+      - 'Makefile'
+      - 'LICENSE'
+      - 'scripts/benchmarks/**'
+  pull_request:
+    branches: [main]
+    paths-ignore:
+      - '**.md'
+      - '**.rst'
+      - 'docs/**'
+      - 'Makefile'
+      - 'LICENSE'
+      - 'scripts/benchmarks/**'
+
+env:
+  UV_FROZEN: "1"
+
+jobs:
+  lint-format:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          enable-cache: true
+          version: "0.12.1"
+      - name: Run lint
+        run: uvx ruff@0.16.1 check --diff
+      - name: Run format
+        run: uvx ruff@0.16.1 format --diff
+
+  tests:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        resolution: ["locked"]
+        # Also test against a fresh resolve of the latest compatible
+        # dependencies so a new upstream release breaks a PR instead of a
+        # release.
+        include:
+          - python-version: "3.13"
+            resolution: "unlocked"
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          python-version: ${{ matrix.python-version }}
+          enable-cache: true
+          version: "0.12.1"
+      - name: Restore Warp kernel cache
+        uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        with:
+          path: ~/.cache/warp
+          key: warp-kernels-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}-${{ matrix.resolution }}-${{ hashFiles('uv.lock', 'mjlab/**/*.py') }}
+          restore-keys: |
+            warp-kernels-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}-${{ matrix.resolution }}-
+      - name: Test with python ${{ matrix.python-version }} (locked deps)
+        if: matrix.resolution == 'locked'
+        run: uv run --extra cpu pytest
+      - name: Test with python ${{ matrix.python-version }} (latest deps)
+        if: matrix.resolution == 'unlocked'
+        env:
+          UV_FROZEN: "0"
+        run: uv run --extra cpu --upgrade pytest
+
+  pyright:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          python-version: ${{ matrix.python-version }}
+          enable-cache: true
+          version: "0.12.1"
+      - name: Test with python ${{ matrix.python-version }}
+        run: uv run --extra cpu pyright
+
+  ty-check:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          python-version: ${{ matrix.python-version }}
+          enable-cache: true
+          version: "0.12.1"
+      - name: Type check with python ${{ matrix.python-version }}
+        run: uv run --extra cpu ty check
+
+  stubs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          enable-cache: true
+          version: "0.12.1"
+      - name: Sync environment
+        run: uv sync --extra cpu
+      - name: Regenerate MuJoCo stubs
+        run: bash typings/generate_mujoco_stubs.sh
+      - name: Verify stubs are up to date
+        run: |
+          if ! git diff --exit-code typings/mujoco; then
+            echo "::error::MuJoCo type stubs are out of date. Run 'make stubs' and commit the result."
+            exit 1
+          fi
+
+  # Mirrors the release smoke test: build the artifacts and import mjlab from an
+  # isolated, freshly resolved install, so packaging and clean-install issues
+  # surface on the PR instead of at publish time.
+  smoke-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          enable-cache: true
+          version: "0.12.1"
+      - name: Install Python 3.13
+        run: uv python install 3.13
+      - name: Build
+        run: uv build
+      # MUJOCO_GL=disable so `import mujoco` skips its GL backend import; the
+      # runner has no GL libraries and the smoke test does not render.
+      - name: Smoke test (wheel)
+        run: uv run --isolated --no-project --with dist/*.whl tests/smoke_test.py
+        env:
+          MUJOCO_GL: disable
+      - name: Smoke test (source distribution)
+        run: uv run --isolated --no-project --with dist/*.tar.gz tests/smoke_test.py
+        env:
+          MUJOCO_GL: disable
--- a/.github/workflows/claude-code-review.yml
+++ b/.github/workflows/claude-code-review.yml
@ -0,0 +1,39 @@
+name: Claude Code Review
+
+on:
+  pull_request:
+    types: [opened, synchronize, ready_for_review, reopened]
+    # Optional: Only run on specific file changes
+    # paths:
+    #   - "src/**/*.ts"
+    #   - "src/**/*.tsx"
+    #   - "src/**/*.js"
+    #   - "src/**/*.jsx"
+
+jobs:
+  claude-review:
+    if: github.event.pull_request.user.login == 'kevinzakka'
+
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code Review
+        id: claude-review
+        uses: anthropics/claude-code-action@be7b93b1907a4abad570368f3c74b6fe3807510b # v1
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+          plugin_marketplaces: 'https://github.com/anthropics/claude-code.git'
+          plugins: 'code-review@claude-code-plugins'
+          prompt: '/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}'
+          # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
+          # or https://code.claude.com/docs/en/cli-reference for available options
--- a/.github/workflows/claude.yml
+++ b/.github/workflows/claude.yml
@ -0,0 +1,49 @@
+name: Claude Code
+
+on:
+  issue_comment:
+    types: [created]
+  pull_request_review_comment:
+    types: [created]
+  issues:
+    types: [opened, assigned]
+  pull_request_review:
+    types: [submitted]
+
+jobs:
+  claude:
+    if: |
+      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
+      (github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+      actions: read # Required for Claude to read CI results on PRs
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code
+        id: claude
+        uses: anthropics/claude-code-action@be7b93b1907a4abad570368f3c74b6fe3807510b # v1
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+
+          # This is an optional setting that allows Claude to read CI results on PRs
+          additional_permissions: |
+            actions: read
+
+          # Optional: Give a custom prompt to Claude. If this is not specified, Claude will perform the instructions specified in the comment that tagged it.
+          # prompt: 'Update the pull request description to include a summary of changes.'
+
+          # Optional: Add claude_args to customize behavior and configuration
+          # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
+          # or https://code.claude.com/docs/en/cli-reference for available options
+          # claude_args: '--allowed-tools Bash(gh pr *)'
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@ -0,0 +1,91 @@
+name: Docker
+
+on:
+  workflow_dispatch:
+
+  push:
+    branches:
+      - "main"
+
+  pull_request:
+    types:
+      - opened
+      - reopened
+      - synchronize
+      - ready_for_review
+
+concurrency:
+  group: docker-${{ github.ref }}
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash
+
+env:
+  FORCE_COLOR: 1
+  REGISTRY: ghcr.io
+  IMAGE_NAME: mujocolab/mjlab
+
+permissions:
+  id-token: write
+  packages: write
+
+jobs:
+  check_paths:
+    runs-on: ubuntu-22.04
+    outputs:
+      build: ${{ steps.filter.outputs.any }}
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - id: filter
+        uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2
+        with:
+          list-files: shell
+          filters: |
+            any:
+              - ".github/workflows/docker.yml"
+              - "Dockerfile"
+
+  build:
+    needs: check_paths
+    if: ${{ needs.check_paths.outputs.build == 'true' }}
+    runs-on: ubuntu-22.04
+
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+
+      - name: Setup Docker buildx
+        uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
+
+      - name: Log into registry
+        uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract Docker metadata
+        id: meta
+        uses: docker/metadata-action@dc802804100637a589fabce1cb79ff13a1411302 # v6.2.0
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+          tags: |
+            type=ref,event=branch
+            type=raw,value=latest,enable={{is_default_branch}}
+
+      - name: Build and push Docker image
+        uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0
+        with:
+          context: .
+          push: ${{ github.ref == 'refs/heads/main' }}
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: |
+            type=gha
+            type=registry,ref=ghcr.io/mujocolab/mjlab/mjlab:buildcache
+          cache-to: |
+            type=gha,mode=max
+            type=registry,ref=ghcr.io/mujocolab/mjlab/mjlab:buildcache,mode=max
+          platforms: linux/amd64
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@ -0,0 +1,47 @@
+name: docs
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - 'v*'
+
+permissions:
+  contents: write
+
+env:
+  UV_FROZEN: "1"
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 0
+
+      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+        with:
+          python-version: '3.13'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          version: "0.12.1"
+
+      - name: Build Sphinx Documentation
+        run: uv run --group docs sphinx-multiversion docs docs/_build
+
+      - name: Add root redirect
+        run: echo '<meta http-equiv="refresh" content="0; url=main/index.html">' > docs/_build/index.html
+
+      - name: Remove Sphinx build artifacts
+        run: find docs/_build -type d -name .doctrees -exec rm -rf {} +
+
+      - name: Deploy to GitHub Pages
+        uses: peaceiris/actions-gh-pages@84c30a85c19949d7eee79c4ff27748b70285e453 # v4.1.0
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          publish_dir: ./docs/_build/
+          keep_files: true
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@ -0,0 +1,37 @@
+name: nightly
+
+# Early warning: upgrade every dependency to its latest version and run the
+# test suite, so an upstream release that breaks mjlab shows up here rather than
+# in a user's install or at release time.
+
+on:
+  schedule:
+    - cron: "17 11 * * *"  # Daily at 11:17 UTC (off the hour to reduce delay).
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  latest-deps:
+    name: Test against latest dependencies (py${{ matrix.python-version }})
+    # Skip on forks; scheduled runs only make sense on the canonical repo.
+    if: github.repository == 'mujocolab/mjlab'
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.13"]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          python-version: ${{ matrix.python-version }}
+          version: "0.12.1"
+      - name: Upgrade all dependencies to latest
+        run: uv lock --upgrade
+      - name: Show dependency tree
+        run: uv tree
+      - name: Test
+        run: uv run --extra cpu pytest
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@ -0,0 +1,40 @@
+name: "Publish"
+
+on:
+  push:
+    tags:
+      - v*
+
+jobs:
+  run:
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - name: Checkout
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Install uv
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          version: "0.12.1"
+      - name: Install Python 3.13
+        run: uv python install 3.13
+      - name: Build
+        run: uv build
+      # MUJOCO_GL=disable so `import mujoco` skips its GL backend import. Without
+      # it, mujoco 3.10 eagerly imports the EGL bindings at import time (mjlab
+      # defaults MUJOCO_GL=egl on Linux) and crashes on the runner, which has no
+      # GL libraries. The smoke test does not render, so disabling GL is fine.
+      - name: Smoke test (wheel)
+        run: uv run --isolated --no-project --with dist/*.whl tests/smoke_test.py
+        env:
+          MUJOCO_GL: disable
+      - name: Smoke test (source distribution)
+        run: uv run --isolated --no-project --with dist/*.tar.gz tests/smoke_test.py
+        env:
+          MUJOCO_GL: disable
+      - name: Publish
+        run: uv publish
--- a/.gitignore
+++ b/.gitignore
@ -0,0 +1,19 @@
+wandb/
+logs/
+onnx/
+videos/
+__pycache__/
+MUJOCO_LOG.TXT
+debug.py
+.vscode/
+*.ipynb_checkpoints/
+motions/
+*_rerun*
+artifacts/
+.venv/
+render_robots.py
+benchmark_results/
+
+# Documentation outputs.
+**/_build/*
+**/generated/*
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@ -0,0 +1,10 @@
+repos:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.14.14
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format
--- a/.python-version
+++ b/.python-version
@ -0,0 +1 @@
+3.13
--- a/AGENTS.md
+++ b/AGENTS.md
@ -0,0 +1 @@
+CLAUDE.md
--- a/CITATION.cff
+++ b/CITATION.cff
@ -0,0 +1,60 @@
+# This CITATION.cff file was generated with cffinit.
+# Visit https://bit.ly/cffinit to generate yours today!
+
+cff-version: 1.2.0
+title: >-
+  mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Kevin
+    family-names: Zakka
+    email: zakka@berkeley.edu
+  - given-names: Brent
+    family-names: Yi
+    email: brentyi@berkeley.edu
+  - given-names: Qiayuan
+    family-names: Liao
+    email: qiayuanl@berkeley.edu
+  - given-names: Louis
+    family-names: Le Lay
+    email: le.lay.louis@gmail.com
+  - given-names: Koushil
+    family-names: Sreenath
+  - given-names: Pieter
+    family-names: Abbeel
+repository-code: 'https://github.com/mujocolab/mjlab'
+keywords:
+  - mujoco
+  - mujoco-warp
+  - simulation
+  - reinforcement-learning
+  - robotics
+license: Apache-2.0
+commit: 0af4087961d6fbe573243951547eb384983aa9c9
+version: 1.6.0
+date-released: '2026-08-08'
+preferred-citation:
+  type: article
+  title: >-
+    mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning
+  authors:
+    - given-names: Kevin
+      family-names: Zakka
+    - given-names: Qiayuan
+      family-names: Liao
+    - given-names: Brent
+      family-names: Yi
+    - given-names: Louis
+      family-names: Le Lay
+    - given-names: Koushil
+      family-names: Sreenath
+    - given-names: Pieter
+      family-names: Abbeel
+  year: 2026
+  url: https://arxiv.org/abs/2601.22074
+  identifiers:
+    - type: arxiv
+      value: 2601.22074
--- a/CLAUDE.md
+++ b/CLAUDE.md
@ -0,0 +1,60 @@
+# Development Workflow
+
+**Always use `uv run`, not python**.
+
+```sh
+
+# 1. Make changes.
+
+# 2. Type check.
+uv run ty check  # Fast
+uv run pyright  # More thorough, but slower
+
+# 3. Run tests.
+uv run pytest tests/  # Single suite
+uv run pytest tests/<test_file>.py  # Specific file
+
+# 4. Format and lint before committing.
+uv run ruff format
+uv run ruff check --fix
+```
+
+We've bundled common commands into a Makefile for convenience.
+
+```sh
+make format     # Format and lint
+make type       # Type-check
+make check      # make format && make type
+make test-fast  # Run tests excluding slow ones
+make test       # Run the full test suite
+make docs       # Build documentation
+```
+
+Always run `make check` before committing. This runs formatting, linting,
+and type checking. Do not commit code that fails type checking.
+
+Before creating a PR, ensure all checks pass with `make test`.
+
+When making user-facing changes, add an entry to `docs/source/changelog.rst`
+under the "Upcoming version (not yet released)" section using
+Added/Changed/Fixed categories. Reference issues with `:issue:\`123\``
+(renders as a link to the GitHub issue).
+
+# Commits and PRs
+
+- Put `Fixes #<number>` at the end of the commit message body, not in
+  the title.
+- PR body should be plain, concise prose. No section headers, checklists,
+  or structured templates. Describe the problem, what the change does, and
+  any non-obvious tradeoffs. A good PR description reads like a short
+  paragraph to a colleague, not a form.
+- PR and commit messages are rendered on GitHub, so don't hard-wrap them
+  at 88 columns. Let each sentence flow on one line.
+
+Some style guidelines to follow:
+- Line length limit is 88 columns. This applies to code, comments, and docstrings.
+- Avoid local imports unless they are strictly necessary (e.g. circular imports).
+- Tests should follow these principles:
+  - Use functions and fixtures; do not use test classes.
+  - Favor targeted, efficient tests over exhaustive edge-case coverage.
+  - Prefer running individual tests rather than the full test suite to improve iteration speed.
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@ -0,0 +1,25 @@
+# Contributing
+
+Bug fixes and documentation improvements are always welcome. For new features, please open an issue first so we can discuss whether it fits and work out the design, as we're intentional about keeping the scope focused.
+
+## Workflow
+
+1. Fork the repository and create a feature branch.
+2. Make your changes.
+3. Ensure formatting, type checking, and tests pass: `make test-all`.
+4. Submit a pull request.
+
+Type checking (`make type`) is required, PRs that don't pass will be blocked. You can optionally install pre-commit hooks (`pre-commit install`) to catch issues early.
+
+## Changelog
+
+Add entries to the "Upcoming version" section in `docs/source/changelog.rst` under the appropriate category (Added / Changed / Fixed), following [Keep a Changelog](https://keepachangelog.com/) conventions.
+
+## Getting Help
+
+- **Issues**: https://github.com/mujocolab/mjlab/issues
+- **Discussions**: https://github.com/mujocolab/mjlab/discussions
+
+## License
+
+By contributing, you agree your contributions will be licensed under Apache 2.0.
--- a/36
+++ b/36
@ -0,0 +1,36 @@
+# Refer to uv-docker-example:
+# https://github.com/astral-sh/uv-docker-example/blob/main/standalone.Dockerfile
+# Note that we use uv to launch, so we omit the second half of the example (non-UV final image)
+
+FROM nvidia/cuda:12.8.0-runtime-ubuntu24.04
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y \
+    git \
+    curl \
+    libegl-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_LINK_MODE=copy
+ENV UV_PYTHON_PREFERENCE=only-managed
+
+RUN uv python install 3.13
+
+WORKDIR /app
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync --locked --no-install-project --no-editable --no-dev
+
+ADD . /app
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --locked --no-editable --no-dev
+
+ENV MUJOCO_GL=egl
+EXPOSE 8080
+
+CMD ["uv", "run", "python", "tests/smoke_test.py"]
--- a/202
+++ b/202
@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2025, The mjlab Developers
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
--- a/74
+++ b/74
@ -0,0 +1,74 @@
+.PHONY: sync
+sync:
+	uv sync --all-packages --extra cu128 --group dev
+
+.PHONY: sync-cpu
+sync-cpu:
+	uv sync --all-packages --extra cpu --group dev
+
+.PHONY: format
+format:
+	uv run ruff format
+	uv run ruff check --fix
+
+.PHONY: type
+type:
+	uv run ty check
+	uv run pyright
+
+.PHONY: stubs
+stubs:
+	bash typings/generate_mujoco_stubs.sh
+
+.PHONY: check
+check: format type
+
+.PHONY: test
+test:
+	uv run pytest
+
+.PHONY: test-fast
+test-fast:
+	uv run pytest -m "not slow"
+
+.PHONY: test-cpu
+test-cpu:
+	FORCE_CPU=1 uv run pytest
+
+.PHONY: test-cpu-fast
+test-cpu-fast:
+	FORCE_CPU=1 uv run pytest -m "not slow"
+
+.PHONY: test-all
+test-all: check test
+
+.PHONY: build
+build:
+	uv build
+	uv run --isolated --no-project --with dist/*.whl tests/smoke_test.py
+	uv run --isolated --no-project --with dist/*.tar.gz tests/smoke_test.py
+	@echo "Build and import test successful"
+
+.PHONY: docs
+docs:
+	uv run --group docs sphinx-build -j auto docs docs/_build
+
+.PHONY: docs-multiversion
+docs-multiversion:
+	uv run --group docs sphinx-multiversion docs docs/_build
+
+.PHONY: docs-watch
+docs-watch:
+	uv run --group docs sphinx-autobuild -j auto docs docs/_build
+
+.PHONY: publish-test
+publish-test: build
+	uv publish --publish-url https://test.pypi.org/legacy/
+
+.PHONY: publish
+publish: build
+	uv publish
+
+.PHONY: docker-build
+docker-build:
+	docker build -t mjlab:latest .
--- a/README.md
+++ b/README.md
@ -0,0 +1,141 @@
+![Project banner](https://raw.githubusercontent.com/mujocolab/mjlab/main/docs/source/_static/mjlab-banner.jpg)
+
+# mjlab
+
+[![GitHub Actions](https://img.shields.io/github/actions/workflow/status/mujocolab/mjlab/ci.yml?branch=main)](https://github.com/mujocolab/mjlab/actions/workflows/ci.yml?query=branch%3Amain)
+[![Documentation](https://github.com/mujocolab/mjlab/actions/workflows/docs.yml/badge.svg)](https://mujocolab.github.io/mjlab/)
+[![License](https://img.shields.io/github/license/mujocolab/mjlab)](https://github.com/mujocolab/mjlab/blob/main/LICENSE)
+[![MuJoCo Warp](https://img.shields.io/badge/MuJoCo_Warp-3.11.0-blue)](https://github.com/google-deepmind/mujoco_warp/releases/tag/v3.11.0)
+[![Nightly Benchmarks](https://img.shields.io/badge/Nightly-Benchmarks-blue)](https://mujocolab.github.io/mjlab/nightly/)
+[![PyPI](https://img.shields.io/pypi/v/mjlab)](https://pypi.org/project/mjlab/)
+[![PyPI downloads](https://img.shields.io/pypi/dm/mjlab?color=blue)](https://pypistats.org/packages/mjlab)
+
+mjlab combines [Isaac Lab](https://github.com/isaac-sim/IsaacLab)'s manager-based API with [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp), a GPU-accelerated version of [MuJoCo](https://github.com/google-deepmind/mujoco).
+The framework provides composable building blocks for environment design,
+with minimal dependencies and direct access to native MuJoCo data structures.
+
+## Getting Started
+
+mjlab requires an NVIDIA GPU for training. macOS is supported for evaluation only.
+
+**Try it now:**
+
+Run the demo (no installation needed):
+
+```bash
+uvx --from mjlab --refresh demo
+```
+
+Or try in [Google Colab](https://colab.research.google.com/github/mujocolab/mjlab/blob/main/notebooks/demo.ipynb) (no local setup required).
+
+**Install from source:**
+
+```bash
+git clone https://github.com/mujocolab/mjlab.git && cd mjlab
+uv run demo
+```
+
+For alternative installation methods (PyPI, Docker), see the [Installation Guide](https://mujocolab.github.io/mjlab/main/source/installation.html).
+
+## Training Examples
+
+### 1. Velocity Tracking
+
+Train a Unitree G1 humanoid to follow velocity commands on flat terrain:
+
+```bash
+uv run train Mjlab-Velocity-Flat-Unitree-G1 --env.scene.num-envs 4096
+```
+
+**Multi-GPU Training:** Scale to multiple GPUs using `--gpu-ids`:
+
+```bash
+uv run train Mjlab-Velocity-Flat-Unitree-G1 \
+  --gpu-ids "[0, 1]" \
+  --env.scene.num-envs 4096
+```
+
+See the [Distributed Training guide](https://mujocolab.github.io/mjlab/main/source/training/distributed_training.html) for details.
+
+Evaluate a policy while training (fetches latest checkpoint from Weights & Biases):
+
+```bash
+uv run play Mjlab-Velocity-Flat-Unitree-G1 --wandb-run-path your-org/mjlab/run-id
+```
+
+### 2. Motion Imitation
+
+Train a humanoid to mimic reference motions. See the [motion imitation guide](https://mujocolab.github.io/mjlab/main/source/training/motion_imitation.html) for preprocessing setup.
+
+```bash
+uv run train Mjlab-Tracking-Flat-Unitree-G1 --registry-name your-org/motions/motion-name --env.scene.num-envs 4096
+uv run play Mjlab-Tracking-Flat-Unitree-G1 --wandb-run-path your-org/mjlab/run-id
+```
+
+### 3. Sanity-check with Dummy Agents
+
+Use built-in agents to sanity check your MDP before training:
+
+```bash
+uv run play Mjlab-Your-Task-Id --agent zero  # Sends zero actions
+uv run play Mjlab-Your-Task-Id --agent random  # Sends uniform random actions
+```
+
+When running motion-tracking tasks, add `--registry-name your-org/motions/motion-name` to the command.
+
+
+## Documentation
+
+Full documentation is available at **[mujocolab.github.io/mjlab](https://mujocolab.github.io/mjlab/)**.
+
+## Development
+
+```bash
+make test          # Run all tests
+make test-fast     # Skip slow tests
+make format        # Format and lint
+make docs          # Build docs locally
+```
+
+For development setup: `uvx pre-commit install`
+
+## Citation
+
+mjlab is used in published research and open-source robotics projects. See the [Research](https://mujocolab.github.io/mjlab/main/source/research.html) page for publications and projects, or share your own in [Show and Tell](https://github.com/mujocolab/mjlab/discussions/categories/show-and-tell).
+
+If you use mjlab in your research, please consider citing:
+
+```bibtex
+@misc{zakka2026mjlablightweightframeworkgpuaccelerated,
+  title={mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning},
+  author={Kevin Zakka and Qiayuan Liao and Brent Yi and Louis Le Lay and Koushil Sreenath and Pieter Abbeel},
+  year={2026},
+  eprint={2601.22074},
+  archivePrefix={arXiv},
+  primaryClass={cs.RO},
+  url={https://arxiv.org/abs/2601.22074},
+}
+```
+
+## License
+
+mjlab is licensed under the [Apache License, Version 2.0](LICENSE).
+
+### Third-Party Code
+
+Some portions of mjlab are forked from external projects:
+
+- **`src/mjlab/utils/lab_api/`** — Utilities forked from [NVIDIA Isaac
+  Lab](https://github.com/isaac-sim/IsaacLab) (BSD-3-Clause license, see file
+  headers)
+
+Forked components retain their original licenses. See file headers for details.
+
+## Acknowledgments
+
+mjlab wouldn't exist without the excellent work of the Isaac Lab team, whose API
+design and abstractions mjlab builds upon.
+
+Thanks to the MuJoCo Warp team — especially Erik Frey and Taylor Howell — for
+answering our questions, giving helpful feedback, and implementing features
+based on our requests countless times.
--- a/RELEASING.md
+++ b/RELEASING.md
@ -0,0 +1,76 @@
+# Releasing
+
+## Pre-release checklist
+
+1. Bump `version` in `pyproject.toml`.
+2. Update `version` and `date-released` in `CITATION.cff`.
+3. Update the "Upcoming version (not yet released)" heading in `docs/source/changelog.rst` to the new version number and date.
+4. Commit the version bump, then create an annotated tag:
+
+```sh
+git tag -a vX.Y.Z -m "Release vX.Y.Z"
+git push origin vX.Y.Z
+```
+
+## Build and verify
+
+Clean previous build artifacts, then build:
+
+```sh
+rm -rf dist/
+make build
+```
+
+This runs `uv build` to produce a wheel and sdist in `dist/`, then smoke-tests
+both artifacts in isolated environments.
+
+## Test on TestPyPI (optional but recommended)
+
+Upload to TestPyPI first to catch packaging issues before the real release:
+
+```sh
+UV_PUBLISH_TOKEN=<your-testpypi-token> make publish-test
+```
+
+Then verify the upload works end-to-end. Use `--index-strategy unsafe-best-match`
+because TestPyPI won't have all dependencies and uv needs to fall back to real
+PyPI for them:
+
+```sh
+uvx --extra-index-url https://test.pypi.org/simple/ \
+    --index-strategy unsafe-best-match \
+    --from mjlab \
+    demo
+```
+
+Note: TestPyPI requires a separate account and token from real PyPI.
+Generate one at https://test.pypi.org/manage/account/token/.
+
+## Publish to PyPI
+
+```sh
+UV_PUBLISH_TOKEN=<your-pypi-token> make publish
+```
+
+Generate a token at https://pypi.org/manage/account/token/.
+
+## Post-release
+
+Verify the release installs and runs correctly. Use `--refresh` to bypass
+the `uvx` cache (which may still hold the TestPyPI version):
+
+```sh
+uvx --refresh --from mjlab demo
+```
+
+## Releasing from a past tag
+
+If the tag has already been created and HEAD has moved ahead, check out the
+tag before building:
+
+```sh
+git checkout vX.Y.Z
+make build
+make publish
+git checkout main
+```
--- a/docs/_templates/versioning.html
+++ b/docs/_templates/versioning.html
@ -0,0 +1,13 @@
+{% if versions %}
+<div class="sidebar-version-switcher">
+  <label class="sidebar-version-label" for="version-select">Version</label>
+  <select id="version-select" class="sidebar-version-select" onchange="location = this.value;">
+    {%- for item in versions.branches %}
+    <option value="{{ item.url }}" {% if item == current_version %}selected{% endif %}>{{ item.name }}</option>
+    {%- endfor %}
+    {%- for item in versions.tags|reverse %}
+    <option value="{{ item.url }}" {% if item == current_version %}selected{% endif %}>{{ item.name }}</option>
+    {%- endfor %}
+  </select>
+</div>
+{% endif %}
--- a/docs/conf.py
+++ b/docs/conf.py
@ -0,0 +1,200 @@
+import os
+import sys
+
+import sphinx_book_theme
+
+sys.path.insert(0, os.path.abspath("../src"))
+sys.path.insert(0, os.path.abspath("../src/mjlab"))
+
+
+project = "mjlab"
+copyright = "2025, The mjlab Developers"
+author = "The mjlab Developers"
+
+extensions = [
+  "sphinx.ext.autodoc",
+  "sphinx.ext.autosummary",
+  "autodocsumm",
+  "myst_parser",
+  "sphinx.ext.napoleon",
+  "sphinxemoji.sphinxemoji",
+  "sphinx.ext.intersphinx",
+  "sphinx.ext.mathjax",
+  "sphinx.ext.todo",
+  "sphinx.ext.viewcode",
+  "sphinxcontrib.bibtex",
+  "sphinxcontrib.icon",
+  "sphinx_copybutton",
+  "sphinx_design",
+  "sphinx_tabs.tabs",
+  "sphinx_multiversion",
+  "sphinx.ext.extlinks",
+]
+
+extlinks = {
+  "issue": (
+    "https://github.com/mujocolab/mjlab/issues/%s",
+    "#%s",
+  ),
+}
+
+mathjax3_config = {
+  "tex": {
+    "inlineMath": [["\\(", "\\)"]],
+    "displayMath": [["\\[", "\\]"]],
+  },
+}
+
+panels_add_bootstrap_css = False
+panels_add_fontawesome_css = True
+
+source_suffix = {
+  ".rst": "restructuredtext",
+  ".md": "markdown",
+}
+
+nitpick_ignore = [
+  ("py:obj", "slice(None)"),
+]
+
+nitpick_ignore_regex = [
+  (r"py:.*", r"pxr.*"),
+  (r"py:.*", r"trimesh.*"),
+]
+
+# emoji style
+sphinxemoji_style = "twemoji"
+autodoc_typehints = "signature"
+autoclass_content = "class"
+autodoc_class_signature = "separated"
+autodoc_member_order = "bysource"
+autodoc_inherit_docstrings = True
+bibtex_bibfiles = ["source/_static/refs.bib"]
+autosummary_generate = True
+autosummary_generate_overwrite = False
+autodoc_default_options = {
+  "member-order": "bysource",
+}
+intersphinx_mapping = {
+  "python": ("https://docs.python.org/3", None),
+}
+
+exclude_patterns = [
+  "_build",
+  "_redirect",
+  "_templates",
+  "Thumbs.db",
+  ".DS_Store",
+  "README.md",
+  "licenses/*",
+]
+
+autodoc_mock_imports = [
+  "matplotlib",
+  "scipy",
+  "carb",
+  "warp",
+  "pxr",
+  "h5py",
+  "hid",
+  "prettytable",
+  "tqdm",
+  "tensordict",
+  "trimesh",
+  "toml",
+  "mjviser",
+  "mujoco_warp",
+  "gymnasium",
+  "rsl_rl",
+  "viser",
+  "wandb",
+  "torchvision",
+]
+
+suppress_warnings = [
+  "ref.python",
+  "docutils",
+]
+
+language = "en"
+
+html_title = "mjlab Documentation"
+html_theme_path = [sphinx_book_theme.get_html_theme_path()]
+html_theme = "sphinx_book_theme"
+html_favicon = "source/_static/favicon.ico"
+html_show_copyright = True
+html_show_sphinx = False
+html_last_updated_fmt = ""
+
+html_static_path = ["source/_static"]
+html_css_files = ["css/custom.css"]
+
+html_theme_options = {
+  "path_to_docs": "docs/",
+  "collapse_navigation": True,
+  "repository_url": "https://github.com/mujocolab/mjlab",
+  "use_repository_button": True,
+  "use_issues_button": True,
+  "use_edit_page_button": True,
+  "show_toc_level": 2,
+  "use_sidenotes": True,
+  "logo": {
+    "text": "mjlab Documentation",
+  },
+  "icon_links": [
+    {
+      "name": "Benchmarks",
+      "url": "https://mujocolab.github.io/mjlab/nightly/",
+      "icon": "fa-solid fa-chart-line",
+      "type": "fontawesome",
+    },
+  ],
+  "icon_links_label": "Quick Links",
+}
+
+templates_path = [
+  "_templates",
+]
+
+smv_remote_whitelist = r"^.*$"
+smv_branch_whitelist = os.getenv("SMV_BRANCH_WHITELIST", r"^(main|devel)$")
+smv_tag_whitelist = os.getenv("SMV_TAG_WHITELIST", r"^v[1-9]\d*\.\d+\.\d+$")
+
+html_sidebars = {
+  "**": [
+    "navbar-logo.html",
+    "search-field.html",
+    "versioning.html",
+    "sbt-sidebar-nav.html",
+  ]
+}
+
+
+def skip_member(app, what, name, obj, skip, options):
+  exclusions = ["from_dict", "to_dict", "replace", "copy", "validate", "__post_init__"]
+  if name in exclusions:
+    return True
+  return None
+
+
+def process_signature(app, what, name, obj, options, signature, return_annotation):
+  """Suppress the ugly __init__ signature for dataclass Cfg classes."""
+  if what == "class" and "exclude-members" in options:
+    if "__init__" in options["exclude-members"]:
+      return ("", None)
+  return None
+
+
+def process_docstring(app, what, name, obj, options, lines):
+  """Strip auto-generated dataclass docstrings (e.g. 'ClassName(*, ...)')."""
+  import dataclasses
+
+  if what == "class" and dataclasses.is_dataclass(obj):
+    if lines and lines[0].startswith(f"{obj.__name__}("):
+      lines.clear()
+
+
+def setup(app):
+  app.connect("autodoc-skip-member", skip_member)
+  app.connect("autodoc-process-signature", process_signature)
+  app.connect("autodoc-process-docstring", process_docstring)
--- a/docs/index.rst
+++ b/docs/index.rst
@ -0,0 +1,128 @@
+Welcome to mjlab!
+=================
+
+.. figure:: source/_static/mjlab-banner.jpg
+   :width: 100%
+   :alt: mjlab
+
+mjlab is a lightweight, open-source framework for robot learning that
+combines GPU-accelerated simulation with composable environments and minimal
+setup friction. It adopts the manager-based API introduced by
+`Isaac Lab <https://github.com/isaac-sim/IsaacLab>`_, where users compose
+modular building blocks for observations, rewards, and events, and pairs it
+with `MuJoCo Warp <https://github.com/google-deepmind/mujoco_warp>`_ for
+GPU-accelerated physics. The result is a framework installable with a single
+command, requiring minimal dependencies, and providing direct access to
+native `MuJoCo <https://github.com/google-deepmind/mujoco>`_ data
+structures.
+
+**Key features:**
+
+- **Composable environments:** users define observations, rewards,
+  terminations, and other MDP terms as modular building blocks
+- **Minimal dependencies:** single-command install via ``uv``, low startup
+  latency
+- **Direct MuJoCo data structures:** native ``MjModel``/``MjData`` access
+  with no translation layers
+- **PyTorch-native:** observations, rewards, and actions are PyTorch
+  tensors backed by zero-copy GPU memory sharing
+
+For more on the design decisions behind mjlab, see :doc:`source/motivation`.
+
+**Try it now** (no installation needed):
+
+.. code-block:: bash
+
+   uvx --from mjlab --refresh demo
+
+Table of Contents
+-----------------
+
+.. toctree::
+   :maxdepth: 1
+   :caption: User Guide
+
+   source/installation
+   source/tutorials
+   source/contributing
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Concepts
+
+   source/architecture_overview
+   source/entity/index
+   source/actuators
+   source/sensors/index
+   source/scene
+   source/terrain
+
+.. toctree::
+   :maxdepth: 1
+   :caption: The Manager Layer
+
+   source/environment_config
+   source/observations
+   source/actions
+   source/rewards
+   source/terminations
+   source/commands
+   source/events
+   source/randomization
+   source/curriculum
+   source/metrics
+   source/recorders
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Training & Debugging
+
+   source/training/rsl_rl
+   source/viewers
+   source/training/distributed_training
+   source/training/cloud
+   source/debugging/nan_guard
+   source/debugging/export_scene
+
+.. toctree::
+   :maxdepth: 2
+   :caption: API Reference
+
+   source/api/index
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Further Reading
+
+   source/motivation
+   source/migration_isaac_lab
+   source/faq
+   source/research
+   source/changelog
+
+License & citation
+------------------
+
+mjlab is licensed under the Apache License, Version 2.0.
+Please refer to the `LICENSE file <https://github.com/mujocolab/mjlab/blob/main/LICENSE/>`_ for details.
+
+If you use mjlab in your research, we would appreciate a citation:
+
+.. code-block:: bibtex
+
+    @article{Zakka_mjlab_A_Lightweight_2026,
+        author = {Zakka, Kevin and Liao, Qiayuan and Yi, Brent and Le Lay, Louis and Sreenath, Koushil and Abbeel, Pieter},
+        title = {{mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning}},
+        url = {https://arxiv.org/abs/2601.22074},
+        year = {2026}
+    }
+
+Acknowledgments
+---------------
+
+mjlab would not exist without the excellent work of the Isaac Lab team, whose API design
+and abstractions mjlab builds upon.
+
+Thanks also to the MuJoCo Warp team — especially Erik Frey and Taylor Howell — for
+answering our questions, giving helpful feedback, and implementing features based
+on our requests countless times.
--- a/docs/source/_static/architecture_diagram.png
+++ b/docs/source/_static/architecture_diagram.png
--- a/docs/source/_static/changelog/mat_texid_dr.gif
+++ b/docs/source/_static/changelog/mat_texid_dr.gif
--- a/docs/source/_static/changelog/native_reward.png
+++ b/docs/source/_static/changelog/native_reward.png
--- a/docs/source/_static/changelog/terrain_visualizer.jpg
+++ b/docs/source/_static/changelog/terrain_visualizer.jpg
--- a/docs/source/_static/content/cartpole-env.jpg
+++ b/docs/source/_static/content/cartpole-env.jpg
--- a/docs/source/_static/content/cartpole_trained.gif
+++ b/docs/source/_static/content/cartpole_trained.gif
--- a/docs/source/_static/content/g1.png
+++ b/docs/source/_static/content/g1.png
--- a/docs/source/_static/content/go1.png
+++ b/docs/source/_static/content/go1.png
--- a/docs/source/_static/content/mjlab-banner.jpg
+++ b/docs/source/_static/content/mjlab-banner.jpg
--- a/docs/source/_static/content/nan_debug.gif
+++ b/docs/source/_static/content/nan_debug.gif
--- a/docs/source/_static/content/rough_terrain.png
+++ b/docs/source/_static/content/rough_terrain.png
--- a/docs/source/_static/content/yam.png
+++ b/docs/source/_static/content/yam.png
--- a/docs/source/_static/css/custom.css
+++ b/docs/source/_static/css/custom.css
@ -0,0 +1,172 @@
+/*
+ * PyData Sphinx Theme — Option A (Indigo/Teal)
+ * Aesthetic: modern lab — indigo primary, teal accent, neutral grays
+ */
+
+/* LIGHT THEME */
+html[data-theme="light"] {
+    /* Brand */
+    --pst-color-primary: #4F46E5;
+    /* Indigo-600 */
+    --pst-color-secondary: #14B8A6;
+    /* Teal-500 */
+    --pst-color-secondary-highlight: #2DD4BF;
+    /* Teal-400 */
+
+    /* Links / code links */
+    --pst-color-inline-code-links: #0D9488;
+    /* Teal-600 */
+    --pst-color-link: var(--pst-color-primary);
+    --pst-color-link-hover: #4338CA;
+    /* Indigo-700 */
+
+    /* Semantic */
+    --pst-color-info: var(--pst-color-secondary);
+    --pst-color-info-highlight: var(--pst-color-secondary);
+    --pst-color-info-bg: #D1FAE5;
+    /* Teal-50 */
+    --pst-color-attention: #F59E0B;
+    /* Amber-500 */
+    --pst-color-target: #EEF2FF;
+    /* Indigo-50 */
+
+    /* Text */
+    --pst-color-text-base: #1F2937;
+    /* Slate-800 */
+    --pst-color-text-muted: #6B7280;
+    /* Slate-500 */
+
+    /* Surfaces */
+    --pst-color-background: #FFFFFF;
+    --pst-color-on-background: #FFFFFF;
+    --pst-color-surface: #F3F4F6;
+    /* Gray-100 */
+    --pst-color-on-surface: #E5E7EB;
+    /* Gray-200 */
+    --pst-color-shadow: #D1D5DB;
+    --pst-color-border: #E5E7EB;
+
+    /* Inline code */
+    --pst-color-inline-code: #0D9488;
+    /* Teal-600 */
+
+    /* Tables / hovers */
+    --pst-color-table-row-hover-bg: #EEF2FF;
+    /* Indigo-50 */
+
+    /* Accent (sparingly) */
+    --pst-color-accent: #10B981;
+    /* Emerald-500 */
+}
+
+/* DARK THEME */
+html[data-theme="dark"] {
+    /* Brand */
+    --pst-color-primary: #A5B4FC;
+    /* Indigo-300/200 mix for readability */
+    --pst-color-secondary: #5EEAD4;
+    /* Teal-300 */
+    --pst-color-secondary-highlight: #2DD4BF;
+
+    /* Links / code links */
+    --pst-color-inline-code-links: #93C5FD;
+    /* Indigo-300 */
+    --pst-color-link: var(--pst-color-primary);
+    --pst-color-link-hover: #818CF8;
+    /* Indigo-400 */
+
+    /* Semantic */
+    --pst-color-info: var(--pst-color-secondary);
+    --pst-color-info-highlight: var(--pst-color-secondary);
+    --pst-color-info-bg: #042F2E;
+    /* Deep teal */
+    --pst-color-attention: #F59E0B;
+    --pst-color-target: #1B1C2A;
+    /* Indigo-tinted surface */
+
+    /* Text */
+    --pst-color-text-base: #E5E7EB;
+    /* Gray-200 */
+    --pst-color-text-muted: #9CA3AF;
+    /* Gray-400 */
+
+    /* Surfaces */
+    --pst-color-background: #0B0C10;
+    /* Deep graphite */
+    --pst-color-on-background: #12131A;
+    --pst-color-surface: #111827;
+    /* Slate-900 */
+    --pst-color-on-surface: #1F2937;
+    /* Slate-800 */
+    --pst-color-shadow: #0F172A;
+    --pst-color-border: #2A2D3A;
+
+    /* Inline code */
+    --pst-color-inline-code: #5EEAD4;
+    /* Teal-300 */
+
+    /* Tables / hovers */
+    --pst-color-table-row-hover-bg: #1B1C2A;
+
+    /* Accent */
+    --pst-color-accent: #34D399;
+    /* Emerald-400 */
+}
+
+/* General tweaks */
+a {
+    text-decoration: none !important;
+}
+
+.bd-header-announcement a,
+.bd-header-version-warning a {
+    color: #5EEAD4;
+}
+
+.form-control {
+    border-radius: 0 !important;
+    border: none !important;
+    outline: none !important;
+}
+
+.navbar-brand,
+.navbar-icon-links {
+    padding-top: 0rem !important;
+    padding-bottom: 0rem !important;
+}
+
+/* Version switcher */
+.sidebar-version-switcher {
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+    padding: 0.4rem 1rem;
+    margin-bottom: 0.5rem;
+}
+
+.sidebar-version-label {
+    font-size: 0.8rem;
+    font-weight: 600;
+    color: var(--pst-color-text-muted);
+    white-space: nowrap;
+}
+
+.sidebar-version-select {
+    flex: 1;
+    font-size: 0.8rem;
+    padding: 0.25rem 0.5rem;
+    border: 1px solid var(--pst-color-border);
+    border-radius: 4px;
+    background: var(--pst-color-background);
+    color: var(--pst-color-text-base);
+    cursor: pointer;
+}
+
+.sidebar-version-select:hover {
+    border-color: var(--pst-color-primary);
+}
+
+/* Sidebar section spacing */
+.bd-sidebar .navbar-icon-links {
+    padding: 0 1rem 0.25rem !important;
+}
--- a/docs/source/_static/dr_combined_rand.gif
+++ b/docs/source/_static/dr_combined_rand.gif
--- a/docs/source/_static/dr_pseudo_inertia.gif
+++ b/docs/source/_static/dr_pseudo_inertia.gif
--- a/docs/source/_static/favicon.ico
+++ b/docs/source/_static/favicon.ico
--- a/docs/source/_static/ghost_visualization.png
+++ b/docs/source/_static/ghost_visualization.png
--- a/docs/source/_static/mjlab-banner.jpg
+++ b/docs/source/_static/mjlab-banner.jpg
--- a/docs/source/_static/native_viewer.png
+++ b/docs/source/_static/native_viewer.png
--- a/docs/source/_static/pattern_grid.jpg
+++ b/docs/source/_static/pattern_grid.jpg
--- a/docs/source/_static/pattern_grid.mp4
+++ b/docs/source/_static/pattern_grid.mp4
--- a/docs/source/_static/pattern_pinhole.mp4
+++ b/docs/source/_static/pattern_pinhole.mp4
--- a/docs/source/_static/ray_alignment_comparison.mp4
+++ b/docs/source/_static/ray_alignment_comparison.mp4
--- a/docs/source/_static/raycast_demo.mp4
+++ b/docs/source/_static/raycast_demo.mp4
--- a/docs/source/_static/refs.bib
+++ b/docs/source/_static/refs.bib
--- a/docs/source/_static/terrains/box_flat.png
+++ b/docs/source/_static/terrains/box_flat.png
--- a/docs/source/_static/terrains/box_inverted_pyramid_stairs.png
+++ b/docs/source/_static/terrains/box_inverted_pyramid_stairs.png
--- a/docs/source/_static/terrains/box_narrow_beams.png
+++ b/docs/source/_static/terrains/box_narrow_beams.png
--- a/docs/source/_static/terrains/box_nested_rings.png
+++ b/docs/source/_static/terrains/box_nested_rings.png
--- a/docs/source/_static/terrains/box_open_stairs.png
+++ b/docs/source/_static/terrains/box_open_stairs.png
--- a/docs/source/_static/terrains/box_pyramid_stairs.png
+++ b/docs/source/_static/terrains/box_pyramid_stairs.png
--- a/docs/source/_static/terrains/box_random_grid.png
+++ b/docs/source/_static/terrains/box_random_grid.png
--- a/docs/source/_static/terrains/box_random_spread.png
+++ b/docs/source/_static/terrains/box_random_spread.png
--- a/docs/source/_static/terrains/box_random_stairs.png
+++ b/docs/source/_static/terrains/box_random_stairs.png
--- a/docs/source/_static/terrains/box_stepping_stones.png
+++ b/docs/source/_static/terrains/box_stepping_stones.png
--- a/docs/source/_static/terrains/box_tilted_grid.png
+++ b/docs/source/_static/terrains/box_tilted_grid.png
--- a/docs/source/_static/terrains/flat_patch_group.png
+++ b/docs/source/_static/terrains/flat_patch_group.png
--- a/docs/source/_static/terrains/hf_discrete_obstacles.png
+++ b/docs/source/_static/terrains/hf_discrete_obstacles.png
--- a/docs/source/_static/terrains/hf_perlin_noise.png
+++ b/docs/source/_static/terrains/hf_perlin_noise.png
--- a/docs/source/_static/terrains/hf_pyramid_slope.png
+++ b/docs/source/_static/terrains/hf_pyramid_slope.png
--- a/docs/source/_static/terrains/hf_random_uniform.png
+++ b/docs/source/_static/terrains/hf_random_uniform.png
--- a/docs/source/_static/terrains/hf_wave.png
+++ b/docs/source/_static/terrains/hf_wave.png
--- a/docs/source/_static/tutorials/cartpole_swingup.mp4
+++ b/docs/source/_static/tutorials/cartpole_swingup.mp4
--- a/docs/source/_static/tutorials/cartpole_training_curve.png
+++ b/docs/source/_static/tutorials/cartpole_training_curve.png
--- a/docs/source/_static/viser_camera_pane.png
+++ b/docs/source/_static/viser_camera_pane.png
--- a/docs/source/_static/viser_viewer.png
+++ b/docs/source/_static/viser_viewer.png
--- a/docs/source/actions.rst
+++ b/docs/source/actions.rst
@ -0,0 +1,173 @@
+.. _actions:
+
+Actions
+=======
+
+Actions define how the policy controls the simulation. The action
+manager receives the policy's output tensor each step, splits it across
+registered action terms, and routes each slice to the appropriate
+entity's actuators. Each term maps a contiguous segment of the policy
+output to a control mode (position, velocity, effort) on a set of
+joints, tendons, or sites.
+
+.. code-block:: python
+
+    from mjlab.envs.mdp.actions import JointPositionActionCfg
+
+    actions = {
+        "joint_pos": JointPositionActionCfg(
+            entity_name="robot",
+            actuator_names=(".*",),   # regex matching actuator names
+            scale=0.5,
+            use_default_offset=True,  # action 0 = default pose
+        ),
+    }
+
+
+Common parameters
+-----------------
+
+All action types share a base set of parameters inherited from
+``BaseActionCfg``.
+
+``entity_name`` identifies the scene entity to control. ``actuator_names``
+is a tuple of regex patterns matched against actuator (or tendon/site)
+names to select the controlled targets.
+
+``scale`` multiplies the raw policy output before any offset is applied.
+It accepts a scalar or a dict mapping actuator name patterns to
+per-target values. This keeps policy outputs in a normalized range while
+mapping to physically meaningful units. ``offset`` is added after
+scaling; joint action types also provide ``use_default_offset``, which
+automatically loads the entity's default joint positions or velocities
+as the offset so that a raw output of zero produces the default pose.
+
+``clip`` optionally clamps the processed action (after scale and offset)
+before it reaches the actuator. It accepts a dict mapping actuator name
+patterns to ``(min, max)`` tuples, resolved the same way as ``scale``
+and ``offset``.
+
+.. code-block:: python
+
+    JointPositionActionCfg(
+        entity_name="robot",
+        actuator_names=(".*",),
+        scale=0.5,
+        clip={".*_hip_.*": (-1.0, 1.0), ".*_knee_.*": (-0.5, 2.0)},
+    )
+
+Actions are written to actuator targets on every decimation substep
+(physics step), not just once per policy step. This is in contrast to
+observation delay, which operates in units of policy steps.
+
+
+Action types
+------------
+
+.. list-table::
+   :header-rows: 1
+   :widths: 28 72
+
+   * - Type
+     - Description
+   * - ``JointPositionAction``
+     - Sets joint position targets. With ``use_default_offset=True``
+       (the default), a policy output of zero commands the default pose.
+       Encoder bias from ``dr.encoder_bias`` is subtracted automatically
+       so that randomized offsets propagate correctly to the control
+       command.
+   * - ``RelativeJointPositionAction``
+     - Sets joint position targets relative to the current joint positions.
+       The target is ``current_pos + action * scale``, so a policy output of
+       zero holds the robot in place regardless of its current configuration.
+   * - ``JointVelocityAction``
+     - Sets joint velocity targets. ``use_default_offset=True`` uses the
+       default joint velocities (typically zero).
+   * - ``JointEffortAction``
+     - Sets joint effort (torque) targets directly. No default offset.
+   * - ``TendonLengthAction``
+     - Sets tendon length targets. Targets are resolved by matching
+       ``actuator_names`` against tendon names.
+   * - ``TendonVelocityAction``
+     - Sets tendon velocity targets.
+   * - ``TendonEffortAction``
+     - Sets tendon effort targets.
+   * - ``SiteEffortAction``
+     - Applies forces and torques at named sites. Useful for
+       quadrotors and drones where thrust is applied at rotor sites
+       rather than through joint actuators.
+
+
+Task-space actions
+------------------
+
+``DifferentialIKAction`` converts Cartesian position and orientation
+commands into joint-space position targets via damped least-squares
+inverse kinematics. One IK step is executed per decimation substep, so
+the end-effector tracks the target continuously across substeps rather
+than only at policy frequency.
+
+The action dimension is selected automatically based on configuration:
+
+- ``orientation_weight == 0``: **3D** (position only)
+- ``orientation_weight > 0, use_relative_mode=True``: **6D** (delta
+  position + delta axis-angle)
+- ``orientation_weight > 0, use_relative_mode=False``: **7D** (absolute
+  position + quaternion)
+
+All objectives (position, orientation, joint limits, posture) are
+stacked into a single DLS system. Setting a weight to zero disables
+that objective with no overhead in the solve.
+
+The ``compute_dq()`` method returns joint displacements without writing
+to actuator targets, enabling multi-iteration IK in standalone scripts
+outside of RL training.
+
+
+Action dimensions and history
+------------------------------
+
+The total action dimension presented to the policy is the sum of each
+registered term's ``action_dim``. For joint, tendon, and site actions
+this equals the number of matched targets. For ``DifferentialIKAction``
+it is 3, 6, or 7 depending on the active objectives.
+
+The action manager tracks the three most recent action vectors:
+``action``, ``prev_action``, and ``prev_prev_action``. Observation terms
+such as ``last_action`` and reward terms such as ``action_rate_l2`` and
+``action_acc_l2`` read from these buffers. Action history is zeroed on
+environment reset so that episode boundaries do not leak information.
+
+
+Multiple action terms
+---------------------
+
+An environment can register any number of terms. The action manager
+concatenates their dimensions in registration order, splits the
+policy's output tensor at the corresponding boundaries, and routes
+each slice independently.
+
+.. code-block:: python
+
+    from mjlab.envs.mdp.actions import (
+        JointPositionActionCfg,
+        JointVelocityActionCfg,
+    )
+
+    actions = {
+        "arm_joints": JointPositionActionCfg(
+            entity_name="robot",
+            actuator_names=(".*_arm_.*",),
+            scale=0.5,
+        ),
+        "wheel_joints": JointVelocityActionCfg(
+            entity_name="robot",
+            actuator_names=(".*_wheel_.*",),
+            scale=10.0,
+        ),
+    }
+
+The policy outputs a tensor whose width equals the total number of
+matched targets across all terms. Terms can also target different
+entities, for example one term for a robot and another for an object
+being manipulated.
--- a/docs/source/actuators.rst
+++ b/docs/source/actuators.rst
@ -0,0 +1,476 @@
+.. _actuators:
+
+Actuators
+=========
+
+Actuators convert high-level commands (position, velocity, effort) into
+low-level efforts that drive joints. They are configured through the
+``articulation`` field of :ref:`EntityCfg <entity>`. mjlab provides
+**built-in** actuators that leverage the physics engine's implicit
+integration for best stability, and **explicit** actuators for custom
+control laws and actuator dynamics.
+
+
+Quick start
+-----------
+
+Basic PD control with ``BuiltinPositionActuator``, the most common
+starting point.
+
+.. code-block:: python
+
+    from mjlab.actuator import BuiltinPositionActuatorCfg
+    from mjlab.entity import EntityCfg, EntityArticulationInfoCfg
+
+    robot_cfg = EntityCfg(
+        spec_fn=lambda: load_robot_spec(),
+        articulation=EntityArticulationInfoCfg(
+            actuators=(
+                BuiltinPositionActuatorCfg(
+                    target_names_expr=(".*_hip_.*", ".*_knee_.*"),
+                    stiffness=80.0,
+                    damping=10.0,
+                    effort_limit=100.0,
+                ),
+            ),
+        ),
+    )
+
+Add delay fields directly on any actuator config to model communication
+latency.
+
+.. code-block:: python
+
+    from mjlab.actuator import BuiltinPositionActuatorCfg
+
+    BuiltinPositionActuatorCfg(
+        target_names_expr=(".*",),
+        stiffness=80.0,
+        damping=10.0,
+        delay_min_lag=2,  # Minimum 2 physics steps
+        delay_max_lag=5,  # Maximum 5 physics steps
+    )
+
+
+Built-in vs explicit actuators
+------------------------------
+
+The key design decision when configuring actuators is whether to use
+**built-in** or **explicit** types. The difference comes down to how
+MuJoCo's integrator handles velocity-dependent forces.
+
+**Built-in actuators** (``BuiltinPositionActuator``,
+``BuiltinVelocityActuator``, ``BuiltinMotorActuator``,
+``BuiltinPdActuator``, ``BuiltinDcMotorActuator``,
+``BuiltinMuscleActuator``) create native MuJoCo actuator elements in the
+MjSpec. The physics engine computes the control law and integrates
+velocity-dependent damping forces implicitly. This provides the best
+numerical stability, particularly with high gains or large timesteps.
+
+**Explicit actuators** (``IdealPdActuator``, ``DcMotorActuator``,
+``LearnedMlpActuator``) compute torques in user code and forward them
+through a ``<motor>`` actuator acting as a passthrough. Because the
+integrator cannot account for the velocity derivatives of these
+externally computed forces, they are less numerically robust than built-in
+types. Use explicit actuators when you need custom control laws or actuator
+dynamics that cannot be expressed with built-in types (e.g.,
+velocity-dependent torque limits, learned actuator networks).
+
+The two approaches match closely in the linear, unconstrained regime at
+small timesteps. At larger timesteps or higher gains, built-in actuators
+are more forgiving.
+
+**Integrator choice.** mjlab places damping inside the actuator rather than
+in joints. The ``euler`` integrator treats joint damping implicitly but
+actuator damping explicitly, limiting stability. The ``implicitfast``
+integrator treats all known velocity-dependent forces implicitly, handling
+both proportional and damping terms of the actuator without additional cost.
+
+.. note::
+
+     mjlab defaults to ``implicitfast``, as it is MuJoCo's recommended
+     integrator and provides superior stability for actuator-side damping.
+
+
+Actuator types
+--------------
+
+All actuator configs share a few common fields inherited from
+``ActuatorCfg``:
+
+- ``target_names_expr``: Tuple of regex patterns matched against joint
+  names (or tendon/site names when using a different
+  ``transmission_type``).
+- ``armature``: Reflected rotor inertia added to the target joint.
+- ``frictionloss``: Static friction (stiction) modeled as a constraint
+  on the target joint. See MuJoCo's
+  `frictionloss <https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-joint-frictionloss>`_.
+
+Built-in actuators
+^^^^^^^^^^^^^^^^^^
+
+Built-in actuators use MuJoCo's native actuator types via the MjSpec API.
+
+**BuiltinPositionActuator**: Creates ``<position>`` actuators for PD
+control.
+
+**BuiltinVelocityActuator**: Creates ``<velocity>`` actuators for velocity
+control.
+
+**BuiltinMotorActuator**: Creates ``<motor>`` actuators for direct torque
+control.
+
+**BuiltinPdActuator**: Native PD that closes on both a position and a
+velocity target, implemented as paired ``<position>`` + ``<velocity>``
+actuators summing to ``kp * (p_target - q) + kd * (v_target - qdot)``.
+``BuiltinPositionActuator`` puts kd on the ``<position>`` element and
+implicitly assumes a zero velocity reference; use this when the policy
+emits a non-zero velocity target. Native delivery lets
+``implicit`` / ``implicitfast`` see the kd term in their velocity update,
+unlike ``IdealPdActuator`` which forwards Python-computed torque through
+an opaque ``<motor>``.
+
+**BuiltinDcMotorActuator**: Wraps MuJoCo's native
+`<dcmotor> <https://mujoco.readthedocs.io/en/stable/XMLreference.html#actuator-dcmotor>`_
+element. Torque is ``tau = K * (V - K * omega) / R``; the back-EMF runs
+through the native bias path, so ``implicit`` / ``implicitfast`` pick up
+its velocity derivative as effective damping. Three input modes pick what
+``ctrl`` carries: VOLTAGE drives the motor directly; POSITION / VELOCITY
+close an internal PID (with anti-windup and slew limiting) against a
+single setpoint, whose Vmax-clamped output becomes torque. POSITION mode
+pins v_target = 0 (the kd term acts on raw velocity). Optional physics:
+inductance,
+thermal model with I^2R heating, cogging ripple, LuGre friction.
+``DcMotorActuator`` (the explicit version) is a software PD with a
+velocity-dependent torque clamp on top of a ``<motor>``; this is the real
+electrical model.
+
+**BuiltinMuscleActuator**: Creates ``<muscle>`` actuators for
+biologically-inspired muscle dynamics with force-length-velocity
+characteristics.
+
+.. code-block:: python
+
+    from mjlab.actuator import BuiltinPositionActuatorCfg, BuiltinVelocityActuatorCfg
+
+    # Mobile manipulator: PD for arm joints, velocity control for wheels.
+    actuators = (
+        BuiltinPositionActuatorCfg(
+            target_names_expr=(".*_shoulder_.*", ".*_elbow_.*", ".*_wrist_.*"),
+            stiffness=100.0,
+            damping=10.0,
+            effort_limit=150.0,
+        ),
+        BuiltinVelocityActuatorCfg(
+            target_names_expr=(".*_wheel_.*",),
+            damping=20.0,
+            effort_limit=50.0,
+        ),
+    )
+
+
+Explicit actuators
+^^^^^^^^^^^^^^^^^^
+
+Explicit actuators compute efforts and forward them to an underlying
+``<motor>`` actuator acting as a passthrough. See
+`Built-in vs explicit actuators`_ above for stability implications.
+
+**IdealPdActuator**: Implements an ideal PD controller. Computes torques
+as ``tau = Kp * pos_error + Kd * vel_error``.
+
+**DcMotorActuator**: Extends ``IdealPdActuator`` with velocity-dependent
+torque saturation to model DC motor torque-speed curves (back-EMF
+effects). Implements a linear torque-speed curve: maximum torque at zero
+velocity, zero torque at maximum velocity.
+
+**LearnedMlpActuator**: Neural network-based actuator that uses a
+trained MLP to predict torque outputs from joint state history. Useful
+when analytical models cannot capture complex actuator dynamics like
+delays, nonlinearities, and friction effects. Inherits DC motor
+velocity-based torque limits.
+
+.. code-block:: python
+
+    from mjlab.actuator import IdealPdActuatorCfg, DcMotorActuatorCfg
+
+    # Ideal PD for hips, DC motor model with torque-speed curve for knees.
+    actuators = (
+        IdealPdActuatorCfg(
+            target_names_expr=(".*_hip_.*",),
+            stiffness=80.0,
+            damping=10.0,
+            effort_limit=100.0,
+        ),
+        DcMotorActuatorCfg(
+            target_names_expr=(".*_knee_.*",),
+            stiffness=80.0,
+            damping=10.0,
+            effort_limit=25.0,       # Continuous torque limit
+            saturation_effort=50.0,  # Peak torque at stall
+            velocity_limit=30.0,     # No-load speed (rad/s)
+        ),
+    )
+
+
+XML actuators
+^^^^^^^^^^^^^
+
+XML actuators wrap actuators already defined in your robot's XML file. The
+config finds existing actuators by matching their ``target`` joint name
+against the ``target_names_expr`` patterns. Each joint must have exactly one
+matching actuator.
+
+**XmlActuator**: Wraps any actuator already defined in the XML. The
+actuator type (position, velocity, motor, muscle) is auto detected from
+the XML element, or you can set ``command_field`` explicitly.
+
+.. code-block:: python
+
+    from mjlab.actuator import XmlActuatorCfg
+
+    # Robot XML already has:
+    # <actuator>
+    #   <position name="hip_joint" joint="hip_joint" kp="100"/>
+    # </actuator>
+
+    # Wrap existing XML actuators.
+    actuators = (
+        XmlActuatorCfg(target_names_expr=("hip_joint",)),
+    )
+
+Actuator delays
+^^^^^^^^^^^^^^^
+
+Any actuator config supports inline delay fields for modeling command
+latency. On a real robot, the onboard PD loop runs at KHz with direct
+encoder access, but the position target from the policy arrives late due
+to inference time and communication bus cycles. Actuator
+delay models this: the command target is delayed, but the control law
+still sees fresh joint state.
+
+This is distinct from observation delay, which models sensor pipeline
+latency (stale state going into the policy). Together they cover both
+legs of the round trip: sensor to policy to motor.
+
+.. code-block:: python
+
+    from mjlab.actuator import IdealPdActuatorCfg
+
+    # Add 2-5 step delay to position commands.
+    actuators = (
+        IdealPdActuatorCfg(
+            target_names_expr=(".*",),
+            stiffness=80.0,
+            damping=10.0,
+            delay_min_lag=2,
+            delay_max_lag=5,
+            delay_hold_prob=0.3,         # 30% chance to keep current lag
+            delay_update_period=10,      # Resample lag every 10 steps
+        ),
+    )
+
+Each step, a lag is sampled uniformly from ``[delay_min_lag,
+delay_max_lag]``. Delays are quantized to physics timesteps. For
+example, with 500Hz physics (2ms/step), ``delay_min_lag=2`` represents
+a 4ms minimum delay.
+
+
+Authoring actuator configs
+--------------------------
+
+Since actuator parameters are uniform within each config, use separate
+actuator configs for joints that need different parameters:
+
+.. code-block:: python
+
+    from mjlab.actuator import BuiltinPositionActuatorCfg
+
+    # G1 humanoid with different gains per joint group.
+    G1_ACTUATORS = (
+        BuiltinPositionActuatorCfg(
+            target_names_expr=(".*_hip_.*", "waist_yaw_joint"),
+            stiffness=180.0,
+            damping=18.0,
+            effort_limit=88.0,
+            armature=0.0015,
+        ),
+        BuiltinPositionActuatorCfg(
+            target_names_expr=("left_hip_pitch_joint", "right_hip_pitch_joint"),
+            stiffness=200.0,
+            damping=20.0,
+            effort_limit=88.0,
+            armature=0.0015,
+        ),
+        BuiltinPositionActuatorCfg(
+            target_names_expr=(".*_knee_joint",),
+            stiffness=150.0,
+            damping=15.0,
+            effort_limit=139.0,
+            armature=0.0025,
+        ),
+        BuiltinPositionActuatorCfg(
+            target_names_expr=(".*_ankle_.*",),
+            stiffness=40.0,
+            damping=5.0,
+            effort_limit=25.0,
+            armature=0.0008,
+        ),
+    )
+
+This design choice reflects a deliberate simplification in mjlab: each
+``ActuatorCfg`` represents a single actuator type (e.g., a specific
+motor/gearbox model) applied uniformly across all joints it drives.
+Hardware parameters such as ``armature`` (reflected rotor inertia) and
+``gear`` describe properties of the actuator hardware, even though they
+are implemented in MuJoCo as joint or actuator fields. In other frameworks
+(like Isaac Lab), these fields may accept ``float | dict[str, float]`` to
+support per-joint variation. mjlab instead encourages one config per
+actuator type or per joint group, keeping the hardware model physically
+consistent and explicit. The main trade-off is verbosity in special cases,
+such as parallel linkages, where per-joint overrides could have been
+convenient, but the benefit is clearer semantics and simpler maintenance.
+
+See :ref:`actions` for how action terms route policy outputs to actuators
+(including DifferentialIK for task-space control), and
+:ref:`domain_randomization` for randomizing gains and effort limits.
+
+
+Computing hardware parameters
+------------------------------
+
+This section is relevant when configuring actuators from real motor
+datasheets. If you are using manually tuned gains, you can skip ahead.
+
+mjlab provides utilities in ``mjlab.utils.actuator`` to compute actuator
+parameters from physical motor specifications. This is particularly
+useful for computing reflected inertia (``armature``) and deriving
+appropriate control gains from hardware datasheets.
+
+**Example: Unitree G1 motor configuration**
+
+.. code-block:: python
+
+    from math import pi
+
+    from mjlab.utils.actuator import (
+        reflected_inertia_from_two_stage_planetary,
+        ElectricActuator
+    )
+
+    # Motor specs from manufacturer datasheet.
+    ROTOR_INERTIAS_7520_14 = (
+        0.489e-4,  # Motor rotor inertia (kg*m**2)
+        0.098e-4,  # Planet carrier inertia
+        0.533e-4,  # Output stage inertia
+    )
+    GEARS_7520_14 = (
+        1,            # First stage (motor to planet)
+        4.5,          # Second stage (planet to carrier)
+        1 + (48/22),  # Third stage (carrier to output)
+    )
+
+    # Compute reflected inertia at joint output.
+    # J_reflected = J_motor*(N1*N2)**2 + J_carrier*N2**2 + J_output.
+    ARMATURE_7520_14 = reflected_inertia_from_two_stage_planetary(
+        ROTOR_INERTIAS_7520_14, GEARS_7520_14
+    )
+
+    # Create motor spec container.
+    ACTUATOR_7520_14 = ElectricActuator(
+        reflected_inertia=ARMATURE_7520_14,
+        velocity_limit=32.0,   # rad/s at joint
+        effort_limit=88.0,     # N*m continuous torque
+    )
+
+    # Derive PD gains from natural frequency and damping ratio.
+    NATURAL_FREQ = 10 * 2*pi  # 10 Hz bandwidth.
+    DAMPING_RATIO = 2.0       # Overdamped, see note below.
+    STIFFNESS = ARMATURE_7520_14 * NATURAL_FREQ**2
+    DAMPING = 2 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ
+
+    # Use in actuator config.
+    from mjlab.actuator import BuiltinPositionActuatorCfg
+
+    actuator = BuiltinPositionActuatorCfg(
+        target_names_expr=(".*_hip_pitch_joint",),
+        stiffness=STIFFNESS,
+        damping=DAMPING,
+        effort_limit=ACTUATOR_7520_14.effort_limit,
+        armature=ACTUATOR_7520_14.reflected_inertia,
+    )
+
+.. note::
+
+     The example uses ``DAMPING_RATIO = 2.0``
+     (overdamped) rather than the critically damped value of 1.0. This is
+     because the reflected inertia calculation only accounts for the motor's
+     rotor inertia, not the apparent inertia of the links being moved. In
+     practice, the total effective inertia at the joint is higher than just
+     the reflected motor inertia, so using an overdamped ratio provides
+     better stability margins when the true system inertia is
+     underestimated.
+
+**Parallel linkage approximation:**
+
+For joints driven by parallel linkages (like the G1's ankles with dual
+motors), the effective armature in the nominal configuration can be
+approximated as the sum of the individual motor armatures:
+
+.. code-block:: python
+
+    # Two 5020 motors driving ankle through parallel linkage.
+    G1_ACTUATOR_ANKLE = BuiltinPositionActuatorCfg(
+        target_names_expr=(".*_ankle_pitch_joint", ".*_ankle_roll_joint"),
+        stiffness=STIFFNESS_5020 * 2,
+        damping=DAMPING_5020 * 2,
+        effort_limit=ACTUATOR_5020.effort_limit * 2,
+        armature=ACTUATOR_5020.reflected_inertia * 2,
+    )
+
+
+Extending: custom actuators
+----------------------------
+
+All actuators implement a unified ``compute()`` interface that receives an
+``ActuatorCmd`` (containing position, velocity, and effort targets) and
+returns control signals for the low-level MuJoCo actuators driving each
+joint.
+
+**Core interface:**
+
+.. code-block:: python
+
+    def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
+        """Convert high-level commands to control signals.
+
+        Args:
+            cmd: Command containing position_target, velocity_target,
+                effort_target (each is a [num_envs, num_targets] tensor
+                or None)
+
+        Returns:
+            Control signals for this actuator
+            ([num_envs, num_targets] tensor)
+        """
+
+**Lifecycle hooks:**
+
+- ``edit_spec``: Modify MjSpec before compilation (add actuators, set
+  gains)
+- ``initialize``: Post-compilation setup (resolve indices, allocate
+  buffers)
+- ``reset``: Per-environment reset logic
+- ``update``: Pre-step updates
+- ``compute``: Convert commands to control signals
+
+**Properties:**
+
+- ``target_ids``: Tensor of local target indices controlled by this
+  actuator
+- ``target_names``: List of target names controlled by this actuator
+- ``ctrl_ids``: Tensor of global control input indices for this actuator
+
+``IdealPdActuator`` is the recommended base class for custom explicit
+actuators. ``DcMotorActuator`` and ``LearnedMlpActuator`` are both
+built on top of it and serve as examples of the extension pattern.
--- a/docs/source/api/actuator.rst
+++ b/docs/source/api/actuator.rst
@ -0,0 +1,182 @@
+mjlab.actuator
+==============
+
+.. automodule:: mjlab.actuator
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`Actuator`
+  - :class:`ActuatorCfg`
+  - :class:`ActuatorCmd`
+  - :class:`BuiltinActuatorGroup`
+  - :class:`BuiltinMotorActuator`
+  - :class:`BuiltinMotorActuatorCfg`
+  - :class:`BuiltinPositionActuator`
+  - :class:`BuiltinPositionActuatorCfg`
+  - :class:`BuiltinVelocityActuator`
+  - :class:`BuiltinVelocityActuatorCfg`
+  - :class:`BuiltinPdActuator`
+  - :class:`BuiltinPdActuatorCfg`
+  - :class:`BuiltinDcMotorActuator`
+  - :class:`BuiltinDcMotorActuatorCfg`
+  - :class:`DcMotorInputMode`
+  - :class:`DcMotorDatasheetParams`
+  - :class:`DcMotorPhysicalParams`
+  - :class:`BuiltinMuscleActuator`
+  - :class:`BuiltinMuscleActuatorCfg`
+  - :class:`XmlActuator`
+  - :class:`XmlActuatorCfg`
+  - :class:`IdealPdActuator`
+  - :class:`IdealPdActuatorCfg`
+  - :class:`DcMotorActuator`
+  - :class:`DcMotorActuatorCfg`
+  - :class:`LearnedMlpActuator`
+  - :class:`LearnedMlpActuatorCfg`
+
+Base
+----
+
+.. autoclass:: Actuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: ActuatorCmd
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Builtin Actuators
+-----------------
+
+.. autoclass:: BuiltinActuatorGroup
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinMotorActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinMotorActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: BuiltinPositionActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinPositionActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: BuiltinVelocityActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinVelocityActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: BuiltinPdActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinPdActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: BuiltinDcMotorActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinDcMotorActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+.. autoclass:: DcMotorInputMode
+  :members:
+  :show-inheritance:
+
+.. autoclass:: DcMotorDatasheetParams
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+.. autoclass:: DcMotorPhysicalParams
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: BuiltinMuscleActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinMuscleActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+XML Actuators
+-------------
+
+.. autoclass:: XmlActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: XmlActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Ideal PD Actuator
+-----------------
+
+.. autoclass:: IdealPdActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: IdealPdActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+DC Motor Actuator
+-----------------
+
+.. autoclass:: DcMotorActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: DcMotorActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Learned MLP Actuator
+--------------------
+
+.. autoclass:: LearnedMlpActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: LearnedMlpActuatorCfg
+  :members:
+  :exclude-members: __init__
--- a/docs/source/api/entity.rst
+++ b/docs/source/api/entity.rst
@ -0,0 +1,45 @@
+mjlab.entity
+============
+
+.. automodule:: mjlab.entity
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`Entity`
+  - :class:`EntityCfg`
+  - :class:`EntityArticulationInfoCfg`
+  - :class:`EntityIndexing`
+  - :class:`EntityData`
+
+Entity
+------
+
+.. autoclass:: Entity
+  :members:
+  :show-inheritance:
+
+.. autoclass:: EntityCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: EntityArticulationInfoCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+EntityIndexing
+--------------
+
+.. autoclass:: EntityIndexing
+  :members:
+
+EntityData
+----------
+
+.. autoclass:: EntityData
+  :members:
--- a/docs/source/api/envs.rst
+++ b/docs/source/api/envs.rst
@ -0,0 +1,36 @@
+mjlab.envs
+==========
+
+.. automodule:: mjlab.envs
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`ManagerBasedRlEnv`
+  - :class:`ManagerBasedRlEnvCfg`
+  - :data:`VecEnvObs`
+  - :data:`VecEnvStepReturn`
+
+ManagerBasedRlEnv
+-----------------
+
+.. autoclass:: ManagerBasedRlEnv
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ManagerBasedRlEnvCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+VecEnvObs
+---------
+
+.. autodata:: VecEnvObs
+
+VecEnvStepReturn
+----------------
+
+.. autodata:: VecEnvStepReturn
--- a/docs/source/api/index.rst
+++ b/docs/source/api/index.rst
@ -0,0 +1,19 @@
+API Reference
+=============
+
+This section provides detailed API documentation for all public modules in mjlab.
+
+.. toctree::
+   :maxdepth: 1
+
+   envs
+   scene
+   sim
+   entity
+   actuator
+   sensor
+   managers
+   terrains
+   rl
+   viewer
+   tasks
--- a/docs/source/api/managers.rst
+++ b/docs/source/api/managers.rst
@ -0,0 +1,208 @@
+mjlab.managers
+==============
+
+.. automodule:: mjlab.managers
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`ManagerBase`
+  - :class:`ManagerTermBase`
+  - :class:`ManagerTermBaseCfg`
+  - :class:`SceneEntityCfg`
+  - :class:`ActionManager`
+  - :class:`ActionTerm`
+  - :class:`ActionTermCfg`
+  - :class:`ObservationManager`
+  - :class:`ObservationGroupCfg`
+  - :class:`ObservationTermCfg`
+  - :class:`RewardManager`
+  - :class:`RewardTermCfg`
+  - :class:`TerminationManager`
+  - :class:`TerminationTermCfg`
+  - :class:`CommandManager`
+  - :class:`NullCommandManager`
+  - :class:`CommandTerm`
+  - :class:`CommandTermCfg`
+  - :class:`CurriculumManager`
+  - :class:`NullCurriculumManager`
+  - :class:`CurriculumTermCfg`
+  - :class:`EventManager`
+  - :class:`EventMode`
+  - :class:`EventTermCfg`
+  - :class:`MetricsManager`
+  - :class:`NullMetricsManager`
+  - :class:`MetricsTermCfg`
+  - :class:`RecorderManager`
+  - :class:`NullRecorderManager`
+  - :class:`RecorderTerm`
+  - :class:`RecorderTermCfg`
+
+Base
+----
+
+.. autoclass:: ManagerBase
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ManagerTermBase
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ManagerTermBaseCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: SceneEntityCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Action Manager
+--------------
+
+.. autoclass:: ActionManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ActionTerm
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ActionTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Observation Manager
+-------------------
+
+.. autoclass:: ObservationManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ObservationGroupCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: ObservationTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Reward Manager
+--------------
+
+.. autoclass:: RewardManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: RewardTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Termination Manager
+-------------------
+
+.. autoclass:: TerminationManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: TerminationTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Command Manager
+---------------
+
+.. autoclass:: CommandManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: NullCommandManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: CommandTerm
+  :members:
+  :show-inheritance:
+
+.. autoclass:: CommandTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Curriculum Manager
+------------------
+
+.. autoclass:: CurriculumManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: NullCurriculumManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: CurriculumTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Event Manager
+-------------
+
+.. autoclass:: EventManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: EventMode
+  :members:
+  :undoc-members:
+
+.. autoclass:: EventTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Metrics Manager
+---------------
+
+.. autoclass:: MetricsManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: NullMetricsManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: MetricsTermCfg
+  :members:
+  :exclude-members: __init__
+
+Recorder Manager
+----------------
+
+.. autoclass:: RecorderManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: NullRecorderManager
+  :members:
+  :show-inheritance:
+
+.. autoclass:: RecorderTerm
+  :members:
+  :show-inheritance:
+
+.. autoclass:: RecorderTermCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
--- a/docs/source/api/rl.rst
+++ b/docs/source/api/rl.rst
@ -0,0 +1,52 @@
+mjlab.rl
+========
+
+.. automodule:: mjlab.rl
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`MjlabOnPolicyRunner`
+  - :class:`RslRlVecEnvWrapper`
+  - :class:`RslRlOnPolicyRunnerCfg`
+  - :class:`RslRlPpoAlgorithmCfg`
+  - :class:`RslRlModelCfg`
+  - :class:`RslRlBaseRunnerCfg`
+
+Runner
+------
+
+.. autoclass:: MjlabOnPolicyRunner
+  :members:
+  :show-inheritance:
+
+.. autoclass:: RslRlVecEnvWrapper
+  :members:
+  :show-inheritance:
+
+Configuration
+-------------
+
+.. autoclass:: RslRlOnPolicyRunnerCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: RslRlPpoAlgorithmCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: RslRlModelCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: RslRlBaseRunnerCfg
+  :members:
+  :exclude-members: __init__
--- a/docs/source/api/scene.rst
+++ b/docs/source/api/scene.rst
@ -0,0 +1,23 @@
+mjlab.scene
+===========
+
+.. automodule:: mjlab.scene
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`Scene`
+  - :class:`SceneCfg`
+
+Scene
+-----
+
+.. autoclass:: Scene
+  :members:
+
+.. autoclass:: SceneCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
--- a/docs/source/api/sensor.rst
+++ b/docs/source/api/sensor.rst
@ -0,0 +1,126 @@
+mjlab.sensor
+============
+
+.. automodule:: mjlab.sensor
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`Sensor`
+  - :class:`SensorCfg`
+  - :class:`SensorContext`
+  - :class:`BuiltinSensor`
+  - :class:`BuiltinSensorCfg`
+  - :class:`ObjRef`
+  - :class:`ContactSensor`
+  - :class:`ContactSensorCfg`
+  - :class:`ContactData`
+  - :class:`ContactMatch`
+  - :class:`RayCastSensor`
+  - :class:`RayCastSensorCfg`
+  - :class:`RayCastData`
+  - :class:`GridPatternCfg`
+  - :class:`PinholeCameraPatternCfg`
+  - :class:`CameraSensor`
+  - :class:`CameraSensorCfg`
+  - :class:`CameraSensorData`
+
+Base
+----
+
+.. autoclass:: Sensor
+  :members:
+  :show-inheritance:
+
+.. autoclass:: SensorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: SensorContext
+  :members:
+
+Builtin Sensor
+--------------
+
+.. autoclass:: BuiltinSensor
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinSensorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: ObjRef
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Contact Sensor
+--------------
+
+.. autoclass:: ContactSensor
+  :members:
+  :show-inheritance:
+
+.. autoclass:: ContactSensorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: ContactData
+  :members:
+
+.. autoclass:: ContactMatch
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Ray Cast Sensor
+---------------
+
+.. autoclass:: RayCastSensor
+  :members:
+  :show-inheritance:
+
+.. autoclass:: RayCastSensorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: RayCastData
+  :members:
+
+.. autoclass:: GridPatternCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: PinholeCameraPatternCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Camera Sensor
+-------------
+
+.. autoclass:: CameraSensor
+  :members:
+  :show-inheritance:
+
+.. autoclass:: CameraSensorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: CameraSensorData
+  :members:
--- a/docs/source/api/sim.rst
+++ b/docs/source/api/sim.rst
@ -0,0 +1,44 @@
+mjlab.sim
+=========
+
+.. automodule:: mjlab.sim
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`Simulation`
+  - :class:`SimulationCfg`
+  - :class:`MujocoCfg`
+  - :class:`TorchArray`
+  - :class:`WarpBridge`
+
+Simulation
+----------
+
+.. autoclass:: Simulation
+  :members:
+
+.. autoclass:: SimulationCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: MujocoCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+TorchArray
+----------
+
+.. autoclass:: TorchArray
+  :members:
+
+WarpBridge
+----------
+
+.. autoclass:: WarpBridge
+  :members:
--- a/docs/source/api/tasks.rst
+++ b/docs/source/api/tasks.rst
@ -0,0 +1,25 @@
+mjlab.tasks
+===========
+
+.. automodule:: mjlab.tasks.registry
+
+.. rubric:: Functions
+
+.. hlist::
+  :columns: 3
+
+  - :func:`register_mjlab_task`
+  - :func:`list_tasks`
+  - :func:`load_env_cfg`
+  - :func:`load_rl_cfg`
+  - :func:`load_runner_cls`
+
+.. autofunction:: register_mjlab_task
+
+.. autofunction:: list_tasks
+
+.. autofunction:: load_env_cfg
+
+.. autofunction:: load_rl_cfg
+
+.. autofunction:: load_runner_cls
--- a/docs/source/api/terrains.rst
+++ b/docs/source/api/terrains.rst
@ -0,0 +1,167 @@
+mjlab.terrains
+==============
+
+.. automodule:: mjlab.terrains
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`TerrainEntity`
+  - :class:`TerrainEntityCfg`
+  - :class:`TerrainGenerator`
+  - :class:`TerrainGeneratorCfg`
+  - :class:`SubTerrainCfg`
+  - :class:`FlatPatchSamplingCfg`
+  - :class:`HfDiscreteObstaclesTerrainCfg`
+  - :class:`HfPerlinNoiseTerrainCfg`
+  - :class:`HfPyramidSlopedTerrainCfg`
+  - :class:`HfRandomUniformTerrainCfg`
+  - :class:`HfWaveTerrainCfg`
+  - :class:`BoxFlatTerrainCfg`
+  - :class:`BoxInvertedPyramidStairsTerrainCfg`
+  - :class:`BoxNarrowBeamsTerrainCfg`
+  - :class:`BoxNestedRingsTerrainCfg`
+  - :class:`BoxOpenStairsTerrainCfg`
+  - :class:`BoxPyramidStairsTerrainCfg`
+  - :class:`BoxRandomGridTerrainCfg`
+  - :class:`BoxRandomSpreadTerrainCfg`
+  - :class:`BoxRandomStairsTerrainCfg`
+  - :class:`BoxSteppingStonesTerrainCfg`
+  - :class:`BoxTiltedGridTerrainCfg`
+
+Core
+----
+
+.. autoclass:: TerrainEntity
+  :members:
+  :show-inheritance:
+
+.. autoclass:: TerrainEntityCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: TerrainGenerator
+  :members:
+
+.. autoclass:: TerrainGeneratorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: SubTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: FlatPatchSamplingCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+Heightfield Terrains
+--------------------
+
+.. autoclass:: HfDiscreteObstaclesTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: HfPerlinNoiseTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: HfRandomUniformTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: HfPyramidSlopedTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: HfWaveTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+Primitive (Box) Terrains
+------------------------
+
+.. autoclass:: BoxFlatTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxInvertedPyramidStairsTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxNarrowBeamsTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxNestedRingsTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxOpenStairsTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxPyramidStairsTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxRandomGridTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxRandomSpreadTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxRandomStairsTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxSteppingStonesTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
+
+.. autoclass:: BoxTiltedGridTerrainCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+  :show-inheritance:
--- a/docs/source/api/viewer.rst
+++ b/docs/source/api/viewer.rst
@ -0,0 +1,72 @@
+mjlab.viewer
+============
+
+.. automodule:: mjlab.viewer
+
+.. rubric:: Classes
+
+.. hlist::
+  :columns: 3
+
+  - :class:`ViewerConfig`
+  - :class:`BaseViewer`
+  - :class:`NativeMujocoViewer`
+  - :class:`ViserPlayViewer`
+  - :class:`OffscreenRenderer`
+
+.. rubric:: Protocols
+
+.. hlist::
+  :columns: 3
+
+  - :class:`EnvProtocol`
+  - :class:`PolicyProtocol`
+  - :class:`VerbosityLevel`
+
+ViewerConfig
+------------
+
+.. autoclass:: ViewerConfig
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+BaseViewer
+----------
+
+.. autoclass:: BaseViewer
+  :members:
+  :show-inheritance:
+
+NativeMujocoViewer
+------------------
+
+.. autoclass:: NativeMujocoViewer
+  :members:
+  :show-inheritance:
+
+ViserPlayViewer
+---------------
+
+.. autoclass:: ViserPlayViewer
+  :members:
+  :show-inheritance:
+
+OffscreenRenderer
+-----------------
+
+.. autoclass:: OffscreenRenderer
+  :members:
+  :show-inheritance:
+
+Protocols
+---------
+
+.. autoclass:: EnvProtocol
+  :members:
+
+.. autoclass:: PolicyProtocol
+  :members:
+
+.. autoclass:: VerbosityLevel
+  :members:
--- a/docs/source/architecture_overview.rst
+++ b/docs/source/architecture_overview.rst
@ -0,0 +1,192 @@
+.. _architecture_overview:
+
+Architecture Overview
+=====================
+
+mjlab is organized into two layers: a **simulation layer** that models
+the robot and world, and a **manager layer** that defines the
+reinforcement learning problem on top of it. Understanding this separation
+is the fastest way to build a mental map of the system.
+
+.. figure:: _static/architecture_diagram.png
+   :width: 60%
+   :align: center
+   :alt: mjlab architecture diagram
+
+   Entities are composed into an MjSpec, compiled, and transferred to
+   MuJoCo Warp for GPU simulation. The ManagerBasedRlEnv orchestrates the
+   MDP; RSL-RL handles training.
+
+
+The simulation layer
+--------------------
+
+**Scene pipeline.**
+mjlab constructs scenes by composing entity descriptions into a single
+`MjSpec <https://mujoco.readthedocs.io/en/stable/programming/modeledit.html>`_.
+Each entity starts from an
+`MJCF <https://mujoco.readthedocs.io/en/latest/modeling.html>`_ file
+loaded via ``MjSpec.from_file()``. Users who define everything in XML can
+use this directly. For more control, Python dataclasses can extend or
+override properties on the loaded spec: actuators, collision rules,
+materials, sensors, and initial state. This hybrid approach lets users
+start from existing MuJoCo models and layer on task-specific configuration
+without modifying the original XML. The composed specification is compiled
+into an ``MjModel`` on the CPU, then transferred to the GPU via
+`MuJoCo Warp <https://mujoco.readthedocs.io/en/stable/mjwarp/index.html>`_,
+which is built on `NVIDIA Warp <https://nvidia.github.io/warp/>`_.
+
+**MuJoCo Warp.**
+MuJoCo Warp is a GPU-accelerated backend for MuJoCo. It preserves
+MuJoCo's ``MjModel``/``MjData`` paradigm but adds a leading *world*
+dimension: a single ``MjData`` object holds the state of N independent
+simulation instances in parallel, enabling thousands of environments to
+be stepped simultaneously. Model parameters are shared across all worlds
+by default, and individual fields can be expanded to vary per-world when
+domain randomization requires it. mjlab captures the simulation step as a
+`CUDA graph <https://developer.nvidia.com/blog/cuda-graphs>`_: the kernel
+execution sequence is recorded once and replayed on subsequent calls,
+eliminating CPU-side dispatch overhead.
+
+.. note::
+
+   CUDA graph capture is a one-time cost at environment startup. Per-episode
+   resets and domain randomization events run as regular Python between graph
+   replays and do not break the capture.
+
+**Components.**
+The simulation layer provides four core components, each with its own
+documentation page:
+
+- :ref:`entity`: a robot, a manipulated object, or a static object such
+  as :ref:`terrain <terrain>`, defined by an MJCF description plus
+  optional Python configuration for actuators, collision rules, and
+  initial state.
+- :ref:`actuators`: how entities are controlled. Users can wrap actuators
+  already defined in MJCF or create new ones from Python configuration.
+- :ref:`sensors`: how the world is observed. Includes MuJoCo-native
+  sensors as well as custom sensors like RGB-D cameras and raycasters.
+- :ref:`scene`: scene composition and environment placement.
+
+
+The manager layer
+-----------------
+
+On top of the simulation layer, mjlab adopts the manager-based environment
+design introduced by Isaac Lab. Users define their environment by composing
+small, self-contained *terms* (reward functions, observation computations,
+domain randomization events) and register them with the appropriate manager.
+Each manager handles the lifecycle of its terms: calling them at the right
+point in the simulation loop, aggregating their outputs, and exposing
+diagnostics.
+
+Terms can be plain functions for stateless computations, or classes that
+inherit from ``ManagerTermBase`` when they need to cache expensive setup
+(such as resolving regex patterns to joint indices at initialization) or
+maintain per-episode state through a ``reset()`` hook.
+
+Environments are configured through ``ManagerBasedRlEnvCfg``, a plain
+dataclass that holds term configuration dictionaries for each manager.
+
+.. code-block:: python
+
+    from mjlab.envs import ManagerBasedRlEnvCfg
+
+    cfg = ManagerBasedRlEnvCfg(
+        decimation=4,           # 4 physics steps per policy step
+        episode_length_s=20.0,
+        scene=...,              # SceneCfg: terrain, entities, sensors
+        sim=...,                # SimulationCfg: timestep, solver, integrator
+        observations={...},     # ObservationManager terms
+        actions={...},          # ActionManager terms
+        rewards={...},          # RewardManager terms
+        terminations={...},     # TerminationManager terms
+        events={...},           # EventManager terms (resets, DR)
+        commands={...},         # CommandManager terms (velocity targets, etc.)
+        curriculum={...},       # CurriculumManager terms
+        metrics={...},          # MetricsManager terms
+    )
+
+.. rubric:: The eight managers
+
+- **ObservationManager**: assembles observation groups with configurable
+  processing (clipping, noise, delay, history). Supports asymmetric
+  actor-critic. See :ref:`observations`.
+- **ActionManager**: routes the policy's output tensor to entity actuators,
+  handling scaling and offset. See :ref:`actions`.
+- **RewardManager**: computes a weighted sum of reward terms, scaled by step
+  duration for frequency invariance. See :ref:`rewards`.
+- **TerminationManager**: evaluates stop conditions, distinguishing terminal
+  resets from timeouts. See :ref:`terminations`.
+- **EventManager**: fires terms at lifecycle points (startup, reset,
+  interval). Domain randomization is implemented through event terms.
+  See :ref:`events` and :ref:`domain_randomization`.
+- **CommandManager**: generates and resamples goal signals (velocity
+  targets, pose targets). See :ref:`commands`.
+- **CurriculumManager**: adjusts training conditions based on policy
+  performance. See :ref:`curriculum`.
+- **MetricsManager**: logs custom per-step values as episode averages.
+  See :ref:`metrics`.
+
+For the full configuration reference covering all managers, see
+:ref:`environment_config`.
+
+
+The environment lifecycle
+-------------------------
+
+Each environment instance passes through four phases.
+
+1. **Build.** ``Scene`` composes entity MJCF files via ``MjSpec`` and
+   compiles ``MjModel`` on the CPU. ``Simulation`` uploads the model to the
+   GPU via MuJoCo Warp, allocating a single ``MjData`` with N parallel
+   worlds. CUDA graphs for ``step``, ``forward``, ``reset``, and ``sense``
+   are captured.
+
+2. **Initialize.** Managers are constructed from the term configuration
+   dictionaries. Regex patterns are matched to joint, body, and geom
+   indices. Observation history and delay buffers are allocated. Model
+   fields required by domain randomization terms are expanded from shared
+   to per-world storage, and CUDA graphs are rebuilt to reflect the new
+   layout. Startup events are fired once.
+
+3. **Reset.** Called at the start of training and whenever an environment
+   terminates or times out. The ``EventManager`` fires ``reset`` terms,
+   which return the scene to an initial state with optional randomization.
+   Command targets are resampled. Observation history buffers are cleared.
+
+4. **Step.** The policy action is processed by the ``ActionManager``. The
+   physics simulation advances ``decimation`` times, with actuator commands
+   applied and entity state updated each sub-step. After the decimation
+   loop, the ``TerminationManager`` checks stop conditions and the
+   ``RewardManager`` computes the reward signal. Step and interval events
+   fire if scheduled, acting on the pre-reset state. Any terminated
+   environments are then reset. A single ``forward()`` call refreshes
+   derived quantities for all environments. The ``CommandManager`` advances
+   or resamples goals. Sensors update. The ``ObservationManager`` assembles
+   the observation for the next policy query.
+
+The step sequence in order:
+
+.. code-block:: text
+
+    action_manager.process_action(action)
+    for _ in range(decimation):
+        action_manager.apply_action()
+        sim.step()
+        scene.update()
+    termination_manager.compute()
+    reward_manager.compute()
+    metrics_manager.compute()
+    event_manager.apply(mode="step")
+    event_manager.apply(mode="interval")
+    [reset terminated envs]
+    sim.forward()
+    command_manager.compute()  # dt=0 for envs just reset
+    sim.sense()
+    observation_manager.compute()
+
+With this mental model in place, the Concepts pages cover each simulation
+layer component in detail, and The Manager Layer pages walk through each
+manager's configuration and built-in terms. If you are coming from Isaac
+Lab, :ref:`migration_isaac_lab` describes the key API differences.
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
--- a/docs/source/commands.rst
+++ b/docs/source/commands.rst
@ -0,0 +1,123 @@
+.. _commands:
+
+Commands
+========
+
+Commands specify what the policy should achieve at each moment: a target
+velocity, a reference trajectory, a goal position. The command manager
+generates these signals, resamples them at configurable intervals, and
+passes them to the policy through the observation system.
+
+
+Registration
+------------
+
+Commands are registered in ``ManagerBasedRlEnvCfg`` as a dictionary
+mapping string names to ``CommandTermCfg`` instances. Unlike the
+function-based terms used by other managers, every command term is a
+class that inherits from ``CommandTerm``.
+
+The ``resampling_time_range`` field controls how often the command
+changes. After each resample the term draws a new timer value uniformly
+from the given ``(min, max)`` range in seconds. Commands are also
+resampled unconditionally on every episode reset.
+
+.. code-block:: python
+
+    commands = {
+        "twist": UniformVelocityCommandCfg(
+            entity_name="robot",
+            resampling_time_range=(3.0, 8.0),
+            ranges=UniformVelocityCommandCfg.Ranges(
+                lin_vel_x=(-1.0, 1.0),
+                lin_vel_y=(-1.0, 1.0),
+                ang_vel_z=(-0.5, 0.5),
+            ),
+        ),
+    }
+
+The ``generated_commands`` observation function reads the current
+command tensor by name and passes it to the policy:
+
+.. code-block:: python
+
+    ObservationTermCfg(
+        func=mdp.generated_commands,
+        params={"command_name": "twist"},
+    )
+
+If the environment has no commands, the manager no-ops all operations
+and returns empty tensors. There is no special handling required.
+
+
+Included command terms
+----------------------
+
+Each task ships with its own command terms tailored to its objective.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 28 72
+
+   * - Term
+     - Description
+   * - ``UniformVelocityCommand``
+     - Generates planar velocity commands ``[v_x, v_y, omega_z]``
+       sampled uniformly from configurable ranges. Supports a standing
+       mode (fraction of environments receive zero velocity) and a
+       heading mode (yaw rate replaced by a proportional controller
+       tracking a sampled heading angle). Used by the velocity task.
+   * - ``LiftingCommand``
+     - Generates a 3D target position for a manipulated object.
+       Supports fixed and dynamic difficulty modes. Tracks metrics
+       including position error and episode success rate. Used by the
+       manipulation task.
+   * - ``MotionCommand``
+     - Streams reference joint positions, velocities, and body poses
+       from a pre-recorded ``.npz`` motion clip. Supports three
+       start-frame sampling modes: ``"start"`` (always frame 0),
+       ``"uniform"`` (random), and ``"adaptive"`` (biased toward
+       difficult regions). At reset the robot is initialized from the
+       sampled frame with optional perturbations. Used by the tracking
+       task.
+
+Each term can render debug visualizations in the interactive viewer
+when ``debug_vis=True`` is set in the configuration. The image below
+shows the ghost visualization from ``MotionCommand``, which renders a
+translucent copy of the robot at the reference pose alongside the
+actual robot.
+
+.. figure:: _static/ghost_visualization.png
+   :align: center
+   :width: 100%
+
+   Viser visualization of the commanded reference motion for the G1 tracking task.
+
+
+Writing custom command terms
+-----------------------------
+
+A custom command term is a class inheriting from ``CommandTerm`` paired
+with a configuration dataclass inheriting from ``CommandTermCfg``. The
+term must implement four methods: ``_resample_command(env_ids)`` to
+sample new goals, ``_update_command(env_ids)`` for per-step updates,
+``_update_metrics()`` for logging, and a ``command`` property returning
+the current goal tensor. The base class manages the resampling timer
+and reset logic automatically.
+
+``_update_command`` is called in two situations. On every environment
+step it receives ``env_ids=None``, meaning update all environments.
+After a reset it is called again with the ids of the environments that
+were just reset, so their command state is brought up to date before
+observations are computed.
+
+The distinction matters when your update advances state, such as
+incrementing a frame index into a reference motion. Apply such advances
+only to ``env_ids`` (all environments when ``None``); otherwise
+resetting a few environments would also advance every other one. Updates
+that simply recompute values from the current simulation state, like a
+heading error, give the same result no matter how often they run and can
+safely ignore ``env_ids``.
+
+The configuration must implement a ``build(env)`` method that
+constructs the paired term instance.
--- a/docs/source/contributing.rst
+++ b/docs/source/contributing.rst
@ -0,0 +1,108 @@
+Contributing
+============
+
+Bug fixes and documentation improvements are always welcome.
+
+.. important::
+
+   For new features, please
+   `open an issue <https://github.com/mujocolab/mjlab/issues>`_ first so
+   we can discuss whether it fits the project scope.
+
+
+Development setup
+-----------------
+
+Clone the repository and sync dependencies:
+
+.. code-block:: bash
+
+   git clone https://github.com/mujocolab/mjlab.git && cd mjlab
+   uv sync
+
+Install pre-commit hooks to catch formatting and lint issues before each
+commit:
+
+.. code-block:: bash
+
+   uvx pre-commit install
+
+
+Common commands
+---------------
+
+The ``Makefile`` provides shortcuts for the most common development tasks:
+
+.. code-block:: bash
+
+   make format      # Format code and fix lint errors (ruff)
+   make type        # Type check (ty + pyright)
+   make check       # Format + type check
+   make test-fast   # Run tests, excluding slow ones
+   make test        # Run the full test suite
+   make test-all    # Format + type check + full test suite
+
+You can also run individual tests for faster iteration:
+
+.. code-block:: bash
+
+   uv run pytest tests/test_rewards.py
+
+Type checking (``make type``) is required. PRs that do not pass will be
+blocked.
+
+
+Building the docs
+-----------------
+
+Build the documentation locally:
+
+.. code-block:: bash
+
+   make docs
+
+The HTML output is written to ``docs/_build/``. For live reload during
+editing:
+
+.. code-block:: bash
+
+   make docs-watch
+
+
+Submitting a pull request
+-------------------------
+
+1. Fork the repository and create a feature branch.
+2. Make your changes.
+3. Run ``make test-all`` to verify formatting, type checking, and tests
+   pass.
+4. Add an entry to the "Upcoming version" section in
+   ``docs/source/changelog.rst`` under the appropriate category
+   (Added / Changed / Fixed), following
+   `Keep a Changelog <https://keepachangelog.com/>`_ conventions.
+5. Submit a pull request.
+
+
+Development with Claude Code
+----------------------------
+
+The repository includes a ``CLAUDE.md`` file at the project root. This file
+defines development conventions, style guidelines, and common commands for
+`Claude Code <https://claude.com/claude-code>`_. It is also a useful
+reference for human contributors since it captures the same rules enforced
+in CI.
+
+The project also includes shared commands in ``.claude/commands/``.
+Any contributor with Claude Code installed can invoke them as slash commands.
+
+``/update-mjwarp <commit-hash>``
+   Update the ``mujoco-warp`` dependency to a specific commit. This edits
+   ``pyproject.toml``, runs ``uv lock``, and opens a PR in one step.
+
+   .. code-block:: text
+
+      /update-mjwarp e28c6038cdf8a353b4146974e4cf37e74dda809a
+
+``/commit-push-pr``
+   Stage current changes, commit, push, and open a PR.
+
--- a/docs/source/curriculum.rst
+++ b/docs/source/curriculum.rst
@ -0,0 +1,184 @@
+.. _curriculum:
+
+Curriculum
+==========
+
+The curriculum manager adjusts training conditions based on policy
+performance. Training begins with an easier problem and difficulty
+increases as the policy demonstrates it can handle the current
+conditions. Common uses include advancing robots to harder terrain,
+widening command velocity ranges, and ramping reward penalty weights
+over the course of training.
+
+Curriculum terms are called at each environment reset. Each term
+receives the environment and the set of resetting environment IDs,
+examines some performance signal, and applies changes to environment
+parameters directly.
+
+.. code-block:: python
+
+    from mjlab.managers.curriculum_manager import CurriculumTermCfg
+
+    curriculum = {
+        "terrain_levels": CurriculumTermCfg(
+            func=mdp.terrain_levels_vel,
+            params={"command_name": "twist"},
+        ),
+    }
+
+The return value of a curriculum function is logged under
+``Curriculum/<term_name>`` in the training metrics.
+
+
+Built-in curriculum functions
+------------------------------
+
+.. list-table::
+   :header-rows: 1
+   :widths: 24 76
+
+   * - Function
+     - Description
+   * - ``terrain_levels_vel``
+     - Measures how far each robot traveled during the episode. Robots
+       that covered enough distance move up one difficulty row in the
+       terrain grid; those that fell short move down. See the terrain
+       curriculum section below.
+   * - ``commands_vel``
+     - Widens velocity command ranges based on training step count.
+       Each stage specifies a step threshold and the new ranges to
+       apply once that threshold is exceeded.
+   * - ``reward_curriculum``
+     - Adjusts a reward term's weight and/or params according to
+       training step thresholds. Replaces the older ``reward_weight``
+       function and also supports modifying reward function parameters.
+   * - ``termination_curriculum``
+     - Adjusts a termination term's params according to training step
+       thresholds. Useful for gradually tightening termination
+       conditions (e.g. energy limits) as training progresses.
+
+
+Reward curriculum
+-----------------
+
+``reward_curriculum`` schedules changes to a reward term's weight or
+keyword arguments as training progresses. Each stage specifies a
+``step`` threshold and an optional ``weight`` or ``params`` update.
+Stages are evaluated in order, and each one whose threshold has been
+reached is applied.
+
+**Ramping a penalty weight**
+
+A common pattern is to introduce a penalty term at low weight early in
+training and increase it once the policy has learned the basics:
+
+.. code-block:: python
+
+    from mjlab.managers.curriculum_manager import CurriculumTermCfg
+
+    curriculum = {
+        "joint_vel_hinge_weight": CurriculumTermCfg(
+            func=mdp.reward_curriculum,
+            params={
+                "reward_name": "joint_vel_hinge",
+                "stages": [
+                    {"step": 0, "weight": -0.01},
+                    {"step": 12000, "weight": -0.1},
+                    {"step": 24000, "weight": -1.0},
+                ],
+            },
+        ),
+    }
+
+**Adjusting reward parameters**
+
+You can also change the parameters passed to the reward function. For
+example, tightening a tracking tolerance as training progresses:
+
+.. code-block:: python
+
+    curriculum = {
+        "track_lin_vel_tighten": CurriculumTermCfg(
+            func=mdp.reward_curriculum,
+            params={
+                "reward_name": "track_linear_velocity",
+                "stages": [
+                    {"step": 0, "params": {"std": 0.5}},
+                    {"step": 20000, "params": {"std": 0.3}},
+                    {"step": 50000, "params": {"std": 0.1}},
+                ],
+            },
+        ),
+    }
+
+**Combining weight and params**
+
+A single stage can update both weight and params at once:
+
+.. code-block:: python
+
+    {"step": 24000, "weight": -1.0, "params": {"max_vel": 1.0}}
+
+
+Termination curriculum
+----------------------
+
+``termination_curriculum`` schedules changes to a termination term's
+parameters as training progresses. This is useful for gradually
+tightening termination conditions once the policy has learned basic
+behaviors.
+
+**Tightening an energy limit**
+
+Start with a permissive energy threshold and reduce it over training:
+
+.. code-block:: python
+
+    from mjlab.managers.curriculum_manager import CurriculumTermCfg
+
+    curriculum = {
+        "energy_threshold": CurriculumTermCfg(
+            func=mdp.termination_curriculum,
+            params={
+                "termination_name": "energy",
+                "stages": [
+                    {"step": 12000, "params": {"threshold": 1000.0}},
+                    {"step": 24000, "params": {"threshold": 700.0}},
+                    {"step": 36000, "params": {"threshold": 400.0}},
+                ],
+            },
+        ),
+    }
+
+The ``time_out`` field on ``TerminationTermCfg`` can also be toggled
+via stages if needed, though this is uncommon in practice.
+
+
+Terrain curriculum
+------------------
+
+The terrain grid used with procedural terrain is a
+``num_rows x num_cols`` matrix of patches. Columns represent terrain
+type variants; rows represent difficulty levels, with row 0 being the
+easiest and row ``num_rows - 1`` the hardest. When
+``TerrainGeneratorCfg.curriculum=True``, each column is assigned exactly
+one terrain type so that difficulty increases monotonically along rows.
+
+At environment construction each environment is assigned a random
+starting row within ``[0, max_init_terrain_level]``. The
+``terrain_levels_vel`` curriculum term promotes or demotes environments
+on each reset based on distance traveled during the episode.
+Environments that reach the maximum level are randomly reassigned to
+any row, maintaining coverage across all difficulty levels. See
+:ref:`terrain` for details on configuring the terrain grid itself.
+
+
+Writing custom curriculum functions
+------------------------------------
+
+A curriculum function accepts ``env`` and ``env_ids``, applies
+parameter changes, and returns a value to log (a scalar tensor, a dict
+of tensors, or ``None``). A typical implementation reads a performance
+metric, decides whether to increase or decrease difficulty, mutates the
+relevant configuration in place, and returns the current difficulty
+level. See :ref:`env-config-term-pattern` for the general pattern.
--- a/docs/source/debugging/export_scene.rst
+++ b/docs/source/debugging/export_scene.rst
@ -0,0 +1,50 @@
+.. _export-scene:
+
+Export Scene
+============
+
+The ``export-scene`` script writes a complete scene (XML and mesh assets) to a
+directory for inspection, sharing, or loading in standalone MuJoCo.
+
+Quick start
+-----------
+
+.. code-block:: bash
+
+    # Export a built-in entity by alias.
+    uv run export-scene g1 --output-dir /tmp/g1
+
+    # Export a registered task scene.
+    uv run export-scene Mjlab-Velocity-Flat-Unitree-Go1 --output-dir /tmp/task
+
+    # Export as a zip archive.
+    uv run export-scene yam --output-dir /tmp/yam --zip True
+
+    # Export a custom entity via import path.
+    uv run export-scene my_pkg.robots:get_my_robot_cfg --output-dir /tmp/custom
+
+The output directory contains a ``scene.xml`` and an ``assets/`` subdirectory
+with all referenced mesh files. The XML can be loaded directly with
+``mujoco.MjModel.from_xml_path()`` or dropped into the
+`simulate viewer <https://mujoco.readthedocs.io/en/stable/programming/samples.html#sasimulate>`_.
+
+Target resolution
+-----------------
+
+The positional ``target`` argument is resolved in order:
+
+1. **Task ID**: checked against the task registry (``import mjlab.tasks``).
+2. **Entity alias**: one of the built-in shorthands (``g1``, ``go1``, ``yam``).
+3. **Import path**: a ``module:attribute`` string pointing to any callable
+   that returns an ``EntityCfg``.
+
+If none match, the script prints available task IDs and aliases.
+
+Options
+-------
+
+``--output-dir DIR`` *(default: "export")*
+    Destination directory. Cleaned before each export to prevent stale assets.
+
+``--zip True`` *(default: False)*
+    Compress the output into a ``.zip`` archive and remove the directory.
--- a/docs/source/debugging/nan_guard.rst
+++ b/docs/source/debugging/nan_guard.rst
@ -0,0 +1,145 @@
+.. _nan-guard:
+
+NaN Guard
+=========
+
+The NaN guard captures simulation states when NaN/Inf is detected, helping
+debug numerical instability issues.
+
+
+Quick start
+-----------
+
+Enable the NaN guard with a single CLI flag:
+
+.. code-block:: bash
+
+    uv run train <task-name> --enable-nan-guard True
+
+This automatically captures and saves simulation states when NaN/Inf is
+detected. You can also enable it programmatically:
+
+.. code-block:: python
+
+    from mjlab.sim.sim import SimulationCfg
+    from mjlab.utils.nan_guard import NanGuardCfg
+
+    cfg = SimulationCfg(
+        nan_guard=NanGuardCfg(
+            enabled=True,
+            buffer_size=100,
+            output_dir="/tmp/mjlab/nan_dumps",
+            max_envs_to_dump=5,
+        ),
+    )
+
+
+Configuration
+-------------
+
+``enabled`` *(default: False)*
+    Enable/disable NaN detection and dumping.
+
+``buffer_size`` *(default: 100)*
+    Number of recent simulation states to keep in the rolling buffer.
+
+``output_dir`` *(default: "/tmp/mjlab/nan_dumps")*
+    Directory where NaN dump files are saved.
+
+``max_envs_to_dump`` *(default: 5)*
+    Maximum number of NaN environments to dump to disk. All environments are
+    tracked in the buffer, but only the first N are saved to reduce dump
+    size.
+
+
+Behavior
+--------
+
+- **Captures** simulation state before each step (``qpos``, ``qvel``,
+  ``act`` if the model has actuator activations, and ``mocap_pos``/``mocap_quat``
+  if the model has mocap bodies)
+- **Detects** NaN/Inf in ``qpos``, ``qvel``, ``qacc``,
+  ``qacc_warmstart``, and ``sensordata`` after each step
+- **Dumps** the rolling buffer and model to disk on first detection
+- **Stops** after the first dump to avoid spam
+
+When disabled, all operations are no-ops with negligible overhead.
+
+
+Output format
+-------------
+
+Each NaN detection creates timestamped files plus latest symlinks:
+
+- ``nan_dump_TIMESTAMP.npz``: compressed state buffer
+
+  - ``states_step_NNNNNN``: captured states per step
+    (shape: ``[num_envs_dumped, state_size]``)
+  - ``_metadata``: dict with ``num_envs_total``, ``nan_env_ids``,
+    ``dumped_env_ids``, etc.
+
+- ``model_TIMESTAMP.mjb``: MuJoCo model in binary format
+- ``nan_dump_latest.npz``: symlink to most recent dump
+- ``model_latest.mjb``: symlink to most recent model
+
+
+Visualizing dumps
+-----------------
+
+Use the interactive viewer to scrub through captured states:
+
+.. code-block:: bash
+
+    # View latest dump.
+    uv run viz-nan /tmp/mjlab/nan_dumps/nan_dump_latest.npz
+
+    # View a specific dump.
+    uv run viz-nan /tmp/mjlab/nan_dumps/nan_dump_20251014_123456.npz
+
+
+.. figure:: ../_static/content/nan_debug.gif
+   :alt: NaN Debug Viewer
+
+   NaN debug viewer.
+
+The viewer provides:
+
+- Step slider to scrub through the buffer
+- Environment slider to compare different environments
+- Info panel showing which environments have NaN/Inf
+- 3D visualization of the robot and terrain at each state
+
+
+NaN detection termination
+-------------------------
+
+While the NaN guard helps **debug** NaN issues by capturing states, you can
+also **prevent** training crashes using the ``nan_detection`` termination
+term. This marks NaN environments as terminated, allowing them to reset
+while training continues:
+
+.. code-block:: python
+
+    from mjlab.envs.mdp.terminations import nan_detection
+    from mjlab.managers.termination_manager import TerminationTermCfg
+
+    nan_term: TerminationTermCfg = field(
+        default_factory=lambda: TerminationTermCfg(
+            func=nan_detection,
+            time_out=False,
+        )
+    )
+
+Terminations are logged as ``Episode_Termination/nan_term`` in your metrics.
+
+.. important::
+
+   ``nan_detection`` is a band-aid, not a cure. If NaNs occur during your
+   task objective (e.g., NaNs happen when grasping), the policy will never
+   learn to complete the task since it resets before receiving rewards.
+   Monitor your ``Episode_Termination/nan_term`` metrics carefully.
+
+**When to use which:**
+
+- ``nan_guard``: debug and understand why NaNs occur (always do this first)
+- ``nan_detection``: keep training stable while working on a permanent fix
--- a/docs/source/entity/entity_data.rst
+++ b/docs/source/entity/entity_data.rst
@ -0,0 +1,507 @@
+.. _entity_data:
+
+Entity Data
+===========
+
+This page is the property reference for ``EntityData``. For an overview
+of how ``entity.data`` fits into the broader data access story, see
+:ref:`entity`.
+
+All properties are PyTorch tensors backed by MuJoCo Warp's GPU buffers
+with no copy overhead. The first dimension is always ``num_envs``, the
+number of parallel simulation worlds.
+
+.. warning::
+
+   Read properties reflect the state after ``sim.forward()`` is called.
+   If you write simulation state and then read a derived property in the
+   same event term, call ``sim.forward()`` between the write and the
+   read. The environment step sequence already does this; the warning
+   applies only when writing custom event terms that mix reads and
+   writes. See the :ref:`FAQ <faq-sim-forward>` for a detailed
+   explanation.
+
+
+Reference: root state
+---------------------
+
+Root properties describe the position, orientation, and velocity of the
+entity's root body. Properties ending in ``_w`` are expressed in the world
+frame. Properties ending in ``_b`` are expressed in the entity's base frame.
+See :ref:`frame-conventions` for details.
+
+Each entity has two root reference points: the **link origin** (the body
+frame origin defined in the MJCF) and the **center of mass (COM)**.
+Which one is relevant depends on the task.
+
+.. admonition:: MuJoCo's mixed-frame ``qvel``
+
+   For floating-base entities, the free joint stores 6 DOFs in
+   ``qvel``. MuJoCo expresses the **linear** components in the
+   **world frame** but the **angular** components in the **local body
+   frame**. EntityData avoids this pitfall: all ``_w`` velocity
+   properties are computed from ``cvel`` (see
+   :ref:`cvel-section` below) and are fully world-frame. If you
+   read ``env.sim.data.qvel`` directly, be aware of the mixed
+   convention.
+
+.. rubric:: Root link properties
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 20 15 30
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``root_link_pose_w``
+     - ``[num_envs, 7]``
+     - world
+     - Root link position (3) and quaternion (4) concatenated
+   * - ``root_link_pos_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root link position
+   * - ``root_link_quat_w``
+     - ``[num_envs, 4]``
+     - world
+     - Root link orientation as quaternion (w, x, y, z)
+   * - ``root_link_vel_w``
+     - ``[num_envs, 6]``
+     - world
+     - Root link linear (3) and angular (3) velocity concatenated
+   * - ``root_link_lin_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root link linear velocity
+   * - ``root_link_ang_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root link angular velocity
+   * - ``root_link_lin_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root link linear velocity in base frame
+   * - ``root_link_ang_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root link angular velocity in base frame
+
+.. rubric:: Root COM properties
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 20 15 30
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``root_com_pose_w``
+     - ``[num_envs, 7]``
+     - world
+     - Root COM position (3) and quaternion (4) concatenated
+   * - ``root_com_pos_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root COM position
+   * - ``root_com_quat_w``
+     - ``[num_envs, 4]``
+     - world
+     - Root COM orientation as quaternion (w, x, y, z)
+   * - ``root_com_vel_w``
+     - ``[num_envs, 6]``
+     - world
+     - Root COM linear (3) and angular (3) velocity concatenated
+   * - ``root_com_lin_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root COM linear velocity
+   * - ``root_com_ang_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root COM angular velocity
+   * - ``root_com_lin_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root COM linear velocity in base frame
+   * - ``root_com_ang_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root COM angular velocity in base frame
+
+.. rubric:: Derived root properties
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 20 15 30
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``projected_gravity_b``
+     - ``[num_envs, 3]``
+     - body
+     - Gravity vector (0, 0, -1) rotated into the base frame. Used to measure
+       tilt: a perfectly upright robot reads ``[0, 0, -1]``.
+   * - ``heading_w``
+     - ``[num_envs]``
+     - world
+     - Heading angle (radians) of the root body's forward axis projected onto
+       the XY plane.
+
+
+Reference: body state
+---------------------
+
+Body properties give per-body kinematic state for all bodies belonging to the
+entity. The second dimension is ``num_bodies``, which counts all non-world
+bodies in the entity's kinematic tree.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 15 25
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``body_link_pose_w``
+     - ``[num_envs, num_bodies, 7]``
+     - world
+     - Per-body link position (3) and quaternion (4)
+   * - ``body_link_pos_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body link positions
+   * - ``body_link_quat_w``
+     - ``[num_envs, num_bodies, 4]``
+     - world
+     - Per-body link orientations
+   * - ``body_link_vel_w``
+     - ``[num_envs, num_bodies, 6]``
+     - world
+     - Per-body link linear (3) and angular (3) velocity
+   * - ``body_link_lin_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body link linear velocities
+   * - ``body_link_ang_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body link angular velocities
+   * - ``body_com_pose_w``
+     - ``[num_envs, num_bodies, 7]``
+     - world
+     - Per-body COM position (3) and quaternion (4)
+   * - ``body_com_pos_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body COM positions
+   * - ``body_com_quat_w``
+     - ``[num_envs, num_bodies, 4]``
+     - world
+     - Per-body COM orientations
+   * - ``body_com_vel_w``
+     - ``[num_envs, num_bodies, 6]``
+     - world
+     - Per-body COM linear (3) and angular (3) velocity
+   * - ``body_com_lin_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body COM linear velocities
+   * - ``body_com_ang_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body COM angular velocities
+   * - ``body_external_wrench``
+     - ``[num_envs, num_bodies, 6]``
+     - world
+     - External force (3) and torque (3) applied to each body
+   * - ``body_external_force``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - External forces applied to each body
+   * - ``body_external_torque``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - External torques applied to each body
+
+
+Reference: joint state
+----------------------
+
+Joint properties cover 1-DOF revolute and prismatic joints. The free joint
+(root floating-base DOF) is excluded; use root state properties for that.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 40
+
+   * - Property
+     - Shape
+     - Description
+   * - ``joint_pos``
+     - ``[num_envs, num_joints]``
+     - Joint positions in radians (revolute) or metres (prismatic)
+   * - ``joint_pos_biased``
+     - ``[num_envs, num_joints]``
+     - Joint positions with encoder bias added. Used when simulating
+       encoder calibration errors via domain randomization.
+   * - ``joint_vel``
+     - ``[num_envs, num_joints]``
+     - Joint velocities in rad/s or m/s
+   * - ``joint_acc``
+     - ``[num_envs, num_joints]``
+     - Joint accelerations in rad/s² or m/s²
+   * - ``actuator_force``
+     - ``[num_envs, num_actuators]``
+     - Scalar actuator output in actuation space (per actuator). This is
+       the force before projection through the transmission Jacobian. For
+       actuator forces in joint space, use ``qfrc_actuator`` instead.
+
+
+.. _generalized-forces:
+
+Reference: generalized forces
+-----------------------------
+
+These properties expose selected components of MuJoCo's generalized
+force decomposition, sliced to this entity's articulated joint DOFs.
+Free joint DOFs are excluded. All shapes are ``[num_envs, nv]`` where
+``nv`` is the number of articulated DOFs belonging to this entity.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Property
+     - Description
+   * - ``qfrc_actuator``
+     - Forces produced by all actuators, mapped into joint space. For
+       motors this is the commanded torque times the gear ratio. For
+       position and velocity actuators this is the force computed by
+       the internal PD law. When ``actuatorgravcomp`` is enabled on a
+       joint, the gravity compensation force is included here.
+   * - ``qfrc_external``
+     - Forces on joints due to Cartesian wrenches applied to bodies
+       via ``xfrc_applied``. This is the :math:`J^\top F` mapping.
+       MuJoCo does not store this term separately; the property
+       recovers it from other force components after ``forward()``.
+
+Reference: geom and site state
+-------------------------------
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 40
+
+   * - Property
+     - Shape
+     - Description
+   * - ``geom_pose_w``
+     - ``[num_envs, num_geoms, 7]``
+     - Per-geom position (3) and quaternion (4) in world frame
+   * - ``geom_pos_w``
+     - ``[num_envs, num_geoms, 3]``
+     - Per-geom positions in world frame
+   * - ``geom_quat_w``
+     - ``[num_envs, num_geoms, 4]``
+     - Per-geom orientations in world frame
+   * - ``geom_vel_w``
+     - ``[num_envs, num_geoms, 6]``
+     - Per-geom linear (3) and angular (3) velocity in world frame
+   * - ``geom_lin_vel_w``
+     - ``[num_envs, num_geoms, 3]``
+     - Per-geom linear velocities in world frame
+   * - ``geom_ang_vel_w``
+     - ``[num_envs, num_geoms, 3]``
+     - Per-geom angular velocities in world frame
+   * - ``site_pose_w``
+     - ``[num_envs, num_sites, 7]``
+     - Per-site position (3) and quaternion (4) in world frame
+   * - ``site_pos_w``
+     - ``[num_envs, num_sites, 3]``
+     - Per-site positions in world frame
+   * - ``site_quat_w``
+     - ``[num_envs, num_sites, 4]``
+     - Per-site orientations in world frame
+   * - ``site_vel_w``
+     - ``[num_envs, num_sites, 6]``
+     - Per-site linear (3) and angular (3) velocity in world frame
+   * - ``site_lin_vel_w``
+     - ``[num_envs, num_sites, 3]``
+     - Per-site linear velocities in world frame
+   * - ``site_ang_vel_w``
+     - ``[num_envs, num_sites, 3]``
+     - Per-site angular velocities in world frame
+
+
+Reference: tendon state
+-----------------------
+
+Tendon properties are only populated for entities that have tendon-driven
+actuators.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 40
+
+   * - Property
+     - Shape
+     - Description
+   * - ``tendon_len``
+     - ``[num_envs, num_tendons]``
+     - Tendon lengths
+   * - ``tendon_vel``
+     - ``[num_envs, num_tendons]``
+     - Tendon velocities
+
+
+.. _frame-conventions:
+
+Frame conventions
+-----------------
+
+Property names encode their reference frame with a suffix.
+
+``_w`` (world frame)
+    A fixed global frame. The origin is typically at the scene origin and
+    its axes are constant throughout the episode. World-frame quantities are
+    useful when you need absolute position, such as checking whether the
+    robot has fallen below a height threshold.
+
+``_b`` (body frame / base frame)
+    The entity's root body frame. It translates and rotates with the robot.
+    Most observation terms use body-frame quantities because they are
+    invariant to the robot's heading direction. A velocity expressed in the
+    body frame reads the same whether the robot faces north or south, which
+    makes it easier for the policy to generalize.
+
+``projected_gravity_b`` is a good example of why the frame suffix
+matters. It takes the world-frame gravity vector ``[0, 0, -1]`` and
+rotates it into the base frame. When the robot is upright the result is
+``[0, 0, -1]``; as the robot tilts, the x and y components grow,
+giving the policy a direct signal for orientation correction.
+
+Quaternion convention
+^^^^^^^^^^^^^^^^^^^^^
+
+All quaternions use the ``(w, x, y, z)`` convention, matching MuJoCo.
+
+Reduced state vs. derived quantities
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+EntityData properties fall into two categories that behave differently
+with respect to ``sim.forward()``:
+
+**Reduced state.** ``joint_pos`` and ``joint_vel`` read directly from
+MuJoCo's ``qpos`` and ``qvel`` arrays. Write methods such as
+``write_joint_state_to_sim()`` modify these arrays directly, so reads
+are always current.
+
+**Derived quantities.** All pose and velocity properties (``*_pose_w``,
+``*_vel_w``, ``*_vel_b``) are computed from MuJoCo's internal arrays
+(``xpos``, ``xquat``, ``cvel``, ``subtree_com``, etc.) which are only
+updated when ``sim.forward()`` runs. If you write to ``qpos``/``qvel``
+and then read a derived property without an intervening ``forward()``,
+the read will return stale values.
+
+The environment step sequence calls ``forward()`` at the right time, so
+this only matters if you write custom event terms that both write and
+read in the same function. See the :ref:`FAQ <faq-sim-forward>` for
+details.
+
+.. _cvel-section:
+
+How velocity properties are computed from ``cvel``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+MuJoCo does not store world-frame linear velocities directly. Instead,
+it stores a 6D spatial velocity per body called ``cvel`` (com-based
+velocity), laid out as ``(angular[3], linear[3])``. This vector is
+expressed in the **c-frame**: a frame centered at ``subtree_com`` (the
+center of mass of the body's kinematic subtree) and oriented like the
+world frame. MuJoCo uses this representation to improve numerical
+precision for mechanisms far from the world origin. See
+`c-frame variables <https://mujoco.readthedocs.io/en/stable/APIreference/APItypes.html#c-frame-variables>`_
+and Featherstone's
+`Spatial Algebra <http://royfeatherstone.org/spatial/>`_ for background.
+
+To recover the world-frame linear velocity at an arbitrary point
+:math:`\mathbf{p}` on a rigid body, we apply the standard rigid-body
+velocity transfer formula. Let :math:`\boldsymbol{\omega}` and
+:math:`\mathbf{v}_c` denote the angular and linear components of
+``cvel``, and let :math:`\mathbf{c}` denote ``subtree_com``. Because
+the c-frame is world-aligned, :math:`\boldsymbol{\omega}` is already in
+the world frame. The linear velocity at :math:`\mathbf{p}` is:
+
+.. math::
+
+   \mathbf{v}_p
+     = \mathbf{v}_c
+       - \boldsymbol{\omega} \times (\mathbf{c} - \mathbf{p})
+
+EntityData applies this formula in ``compute_velocity_from_cvel()``:
+
+.. code-block:: python
+
+   def compute_velocity_from_cvel(pos, subtree_com, cvel):
+       lin_vel_c = cvel[..., 3:6]
+       ang_vel_c = cvel[..., 0:3]
+       offset = subtree_com - pos
+       lin_vel_w = lin_vel_c - torch.cross(ang_vel_c, offset, dim=-1)
+       ang_vel_w = ang_vel_c
+       return torch.cat([lin_vel_w, ang_vel_w], dim=-1)
+
+Every velocity property in EntityData (``root_link_vel_w``,
+``body_link_vel_w``, ``geom_vel_w``, ``site_vel_w``, and their COM
+variants) uses this function, substituting the appropriate point:
+
+- **Link velocities** use ``xpos`` (body frame origin).
+- **COM velocities** use ``xipos`` (body center of mass).
+- **Geom/site velocities** use ``geom_xpos``/``site_xpos``, with
+  ``cvel`` looked up from the parent body.
+
+
+Default pose and relative quantities
+--------------------------------------
+
+``entity.data.default_joint_pos`` holds the joint positions from the entity's
+initial-state configuration (the ``init_state.joint_pos`` field of
+``EntityCfg``). It has shape ``[num_envs, num_joints]`` and is replicated
+across all environments at initialization time.
+
+The relative joint position is the deviation of the current joint position
+from this default:
+
+.. code-block:: python
+
+    joint_pos_rel = joint_pos - default_joint_pos
+
+This is what the ``joint_pos_rel`` observation function computes:
+
+.. code-block:: python
+
+    def joint_pos_rel(env, asset_cfg):
+        asset = env.scene[asset_cfg.name]
+        jnt_ids = asset_cfg.joint_ids
+        return (
+            asset.data.joint_pos[:, jnt_ids]
+            - asset.data.default_joint_pos[:, jnt_ids]
+        )
+
+Relative joint positions give the policy a compact representation of posture
+deviation. When the robot is at its default pose, every element is zero.
+
+Similarly, ``default_joint_vel`` is used by the ``joint_vel_rel`` observation
+function. For most configurations the default velocity is zero, so
+``joint_vel_rel`` is identical to ``joint_vel``. The indirection exists to
+allow non-zero reference velocities in tasks such as motion imitation.
+
+The ``use_default_offset=True`` option in joint position action configs uses
+``default_joint_pos`` as the zero point for the action space, so a network
+output of zero commands the robot to its default pose. This is the standard
+configuration for locomotion tasks.
--- a/docs/source/entity/index.rst
+++ b/docs/source/entity/index.rst
@ -0,0 +1,355 @@
+.. _entity:
+
+Entity
+======
+
+An ``Entity`` represents a physical object in the simulation: a robot, a
+manipulated object, or a fixed fixture like a table. It is the central
+abstraction in mjlab's physics layer.
+
+A single ``Entity`` class covers all variants (contrast Isaac Lab, which
+splits this across ``Articulation``, ``RigidObject``, and several other
+subclasses of ``AssetBase``). Two orthogonal boolean properties classify
+each instance:
+
+**Base type.**
+  A *fixed-base* entity is welded to the world and has no free joint. A
+  *floating-base* entity has a free joint giving it 6-DOF movement.
+
+**Articulation.**
+  An *articulated* entity has internal joints (revolute, prismatic, etc.).
+  A *non-articulated* entity has none beyond a possible free joint.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 25 15 15 15
+
+   * - Type
+     - Example
+     - ``is_fixed_base``
+     - ``is_articulated``
+     - ``is_actuated``
+   * - Fixed non-articulated
+     - Table, wall
+     - True
+     - False
+     - False
+   * - Fixed articulated
+     - Robot arm, door
+     - True
+     - True
+     - True/False
+   * - Floating non-articulated
+     - Box, ball, mug
+     - False
+     - False
+     - False
+   * - Floating articulated
+     - Humanoid, quadruped
+     - False
+     - True
+     - True/False
+
+.. note::
+
+   mjlab automatically wraps every fixed-base entity in a
+   `mocap body <https://mujoco.readthedocs.io/en/stable/modeling.html#mocap-bodies>`_
+   so that each parallel environment can place the entity at a different
+   position. Without this wrapping, all fixed-base entities would be
+   welded to the world origin. The wrapping is transparent, but
+   **positioning only happens when a reset event runs**. You must
+   include a reset event such as ``reset_root_state_uniform`` in your
+   event config; without one, every fixed-base entity will remain at
+   the origin. See the :ref:`FAQ <faq>` for a full example. Mocap
+   entities can also be repositioned at runtime via
+   ``entity.write_mocap_pose_to_sim()``.
+
+
+Configuring an entity
+---------------------
+
+Every entity is described by an ``EntityCfg``. Only ``spec_fn`` is
+required in practice; all other fields have sensible defaults. A passive
+floating object needs nothing more than:
+
+.. code-block:: python
+
+    from mjlab.entity import EntityCfg
+
+    cube_cfg = EntityCfg(spec_fn=get_cube_spec)
+
+An actuated robot uses more of the interface:
+
+.. code-block:: python
+
+    from mjlab.entity import EntityCfg, EntityArticulationInfoCfg
+    from mjlab.actuator import IdealPDActuatorCfg
+
+    robot_cfg = EntityCfg(
+        spec_fn=get_spec,
+        init_state=EntityCfg.InitialStateCfg(
+            pos=(0.0, 0.0, 0.8),
+            joint_pos={".*_hip_.*": 0.5, ".*": 0.0},
+        ),
+        articulation=EntityArticulationInfoCfg(
+            actuators=(
+                IdealPDActuatorCfg(
+                    target_names_expr=(".*",),
+                    stiffness={".*": 50.0},
+                    damping={".*": 5.0},
+                ),
+            ),
+        ),
+        collisions=(my_collision_cfg,),
+    )
+
+The following sections describe each field.
+
+``spec_fn``
+^^^^^^^^^^^
+
+A callable that returns an ``mujoco.MjSpec``. The scene calls it during
+composition, attaches the returned spec with a name prefix, and compiles
+everything into a shared ``MjModel``.
+
+For simple cases a lambda suffices:
+
+.. code-block:: python
+
+    spec_fn = lambda: mujoco.MjSpec.from_file("robot.xml")
+
+For anything more involved, use a regular function. MuJoCo resolves mesh
+assets from disk automatically, so ``get_spec`` only needs to load the
+XML:
+
+.. code-block:: python
+
+    def get_spec() -> mujoco.MjSpec:
+        return mujoco.MjSpec.from_file(str(ROBOT_XML))
+
+Because ``spec_fn`` is an arbitrary callable, you can perform any
+`MjSpec edits <https://mujoco.readthedocs.io/en/stable/python.html#spec>`_
+before returning: add bodies, change joint limits, swap materials,
+or build the entire model programmatically without an XML file at all.
+
+``init_state``
+^^^^^^^^^^^^^^
+
+Default root pose, root velocity, and joint positions/velocities. These
+values are stored as a MuJoCo keyframe and used by reset events to
+return the entity to its initial configuration.
+
+``joint_pos`` and ``joint_vel`` are dicts mapping regex patterns to
+values. Patterns are matched against joint names in order, so later
+entries override earlier ones for any joint that matches both:
+
+.. code-block:: python
+
+    init_state = EntityCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 0.8),       # root position
+        rot=(1.0, 0.0, 0.0, 0.0),  # root quaternion (w, x, y, z)
+        joint_pos={
+            ".*": 0.0,              # all joints to zero
+            ".*_hip_.*": 0.5,       # then override hips to 0.5
+        },
+    )
+
+Set ``joint_pos=None`` to use an existing keyframe from the MJCF model
+instead of defining values here.
+
+``articulation``
+^^^^^^^^^^^^^^^^
+
+Actuator configuration. Only needed for entities that have actuated
+joints. Passive objects (boxes, tables, walls) can omit this field
+entirely. See :ref:`actuators` for details on actuator types.
+
+``soft_joint_pos_limit_factor`` (default 1.0) shrinks the joint range
+used by soft-limit penalty rewards, so the policy is penalized before
+reaching the physical hard stop. This does not modify the actual joint
+limits in the MuJoCo model.
+
+Spec editors
+^^^^^^^^^^^^
+
+The remaining fields are optional tuples of spec editor configs that
+modify the ``MjSpec`` before compilation:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 20 80
+
+   * - Field
+     - Purpose
+   * - ``collisions``
+     - Replace the entity's collision structure: which geoms collide, and
+       with what contact parameters.
+   * - ``lights``
+     - Add lights to specific bodies.
+   * - ``cameras``
+     - Add cameras to specific bodies.
+   * - ``textures``
+     - Add procedural textures (checker, gradient, etc.).
+   * - ``materials``
+     - Add materials and optionally assign them to geoms by regex.
+   * - ``geoms``
+     - Patch attributes of existing geoms (visualization group, collision
+       attributes). Unset attributes are left untouched.
+
+Each editor accepts regex patterns to target specific elements. For
+example, a ``CollisionCfg`` with ``geom_names_expr=(".*_foot.*",)``
+sets contact parameters only on foot geoms. See the asset zoo
+(``mjlab.asset_zoo.robots``) for complete examples.
+
+``geoms`` and ``collisions`` both write geom attributes but with
+different semantics. A ``GeomCfg`` is a sparse *patch*: every attribute
+defaults to ``None``, and only attributes you set are written. A
+``CollisionCfg`` is a *policy*: ``contype``, ``conaffinity``,
+``condim``, and ``priority`` are required and always written to every
+matched geom, and non-matched geoms have collision disabled by default,
+so the entity's contact behavior is fully determined by the config
+regardless of the source XML. Collision configs are applied after geom
+configs; mjlab warns if a ``GeomCfg`` sets a collision attribute that a
+``CollisionCfg`` then overwrites.
+
+Heterogeneous worlds
+^^^^^^^^^^^^^^^^^^^^
+
+For scenes that need different mesh assets in different parallel worlds
+(for example, training a manipulation policy that generalizes across
+object shapes), use ``VariantEntityCfg`` instead of ``EntityCfg``. Each
+world is assigned a variant proportional to a configurable weight, and
+mesh-dependent compiled constants (collision bounds, body inertials,
+subtree mass) are stored as per-world arrays so domain randomization and
+viewers stay consistent. See :ref:`heterogeneous_worlds`.
+
+Subclassing Entity
+^^^^^^^^^^^^^^^^^^
+
+``Entity`` and ``EntityCfg`` can be subclassed for specialized behavior.
+mjlab itself does this for terrain: ``TerrainEntity`` extends ``Entity``
+with procedural terrain generation and per-environment origin
+computation, and ``TerrainEntityCfg`` adds fields like
+``terrain_type``, ``env_spacing``, and ``terrain_generator``. The same
+pattern works for any domain-specific entity that needs logic beyond
+what ``EntityCfg`` and spec editors provide.
+
+Finding elements
+^^^^^^^^^^^^^^^^
+
+Entity provides ``find_*`` methods that accept regex patterns and return
+matched element indices and names:
+
+.. code-block:: python
+
+    ids, names = entity.find_joints((".*_hip_.*", ".*_knee_.*"))
+    ids, names = entity.find_geoms((".*foot.*",))
+    ids, names = entity.find_bodies((".*",))
+
+Available methods: ``find_bodies()``, ``find_joints()``,
+``find_geoms()``, ``find_sites()``, ``find_tendons()``.
+These are used internally during scene construction and manager
+initialization. In reward and observation terms, prefer
+``SceneEntityCfg`` with name patterns as described below.
+
+
+Reading runtime state
+---------------------
+
+Once entities are added to a ``SceneCfg`` and the environment is
+constructed, their state is accessible through three interfaces at
+decreasing levels of abstraction.
+
+EntityData
+^^^^^^^^^^
+
+``entity.data`` is the primary interface for reward, observation, and
+termination functions. It exposes kinematic state (poses, velocities, accelerations), actuator forces,
+generalized forces, and derived body-frame quantities such as projected
+gravity, all as PyTorch tensors with
+shape ``(num_envs, ...)``. See :ref:`entity_data` for the full property
+reference.
+
+``SceneEntityCfg`` selects which entity and which elements within it a
+term operates on. Regex patterns in ``joint_names``, ``body_names``,
+``site_names``, etc. are resolved to integer indices once at manager
+initialization, so there is no regex overhead at runtime:
+
+.. code-block:: python
+
+    from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+    def flat_orientation_l2(
+        env,
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+    ) -> torch.Tensor:
+        """Penalize non-flat base orientation using projected gravity."""
+        asset = env.scene[asset_cfg.name]
+        return torch.sum(
+            torch.square(asset.data.projected_gravity_b[:, :2]), dim=1
+        )
+
+``SceneEntityCfg`` also supports regex element selection through
+``joint_names``, ``body_names``, ``site_names``, etc. The resolved
+integer indices (e.g., ``asset_cfg.joint_ids``) make the runtime read a
+single tensor slice with no regex overhead.
+
+Sensors
+^^^^^^^
+
+Sensors are configured on the **scene**, not on individual entities.
+A sensor can reference an entity element (e.g., a contact sensor on the
+robot's feet, an accelerometer attached to a body site), but it can also
+be independent of any entity. This is why sensors live in ``SceneCfg``
+rather than ``EntityCfg``.
+
+At runtime, sensors are accessed by name through ``env.scene``, the same
+way entities are:
+
+.. code-block:: python
+
+    def angular_momentum_penalty(env, sensor_name: str) -> torch.Tensor:
+        sensor = env.scene[sensor_name]
+        return torch.sum(torch.square(sensor.data), dim=-1)
+
+Builtin sensors wrap MuJoCo sensor types (accelerometer, gyro, framepos,
+subtreeangmom, etc.). ``ContactSensor``, ``RayCastSensor``, and
+``CameraSensor`` provide higher-level abstractions for contact detection,
+terrain scanning, and RGB-D rendering. See :ref:`sensors` for details.
+
+Raw simulation data
+^^^^^^^^^^^^^^^^^^^
+
+For anything not covered by ``EntityData`` or sensors, the underlying
+MuJoCo Warp arrays are accessible through ``env.sim.data`` and
+``env.sim.model``. These expose the full ``mjData`` and ``mjModel``
+fields as PyTorch tensors (zero-copy), indexed by global MuJoCo IDs
+rather than per-entity IDs:
+
+.. code-block:: python
+
+    # Global joint positions across all entities.
+    qpos = env.sim.data.qpos          # (num_envs, nq)
+
+    # All body positions.
+    xpos = env.sim.data.xpos          # (num_envs, nbody, 3)
+
+    # Model-level constants.
+    body_mass = env.sim.model.body_mass  # (nbody,)
+
+This is useful for low-level operations or when you need quantities
+that span multiple entities.
+
+.. note::
+
+   The main limitation of raw sim data is that you must manage global
+   MuJoCo indices yourself. In the future, we plan to support MuJoCo's
+   `bind <https://mujoco.readthedocs.io/en/latest/python.html#relationship-to-pymjcf-and-bind>`_
+   functionality, which will allow binding spec elements directly to
+   their corresponding data views without manual index bookkeeping.
+
+.. toctree::
+   :maxdepth: 1
+
+   entity_data
+   per_world_mesh
--- a/docs/source/entity/per_world_mesh.rst
+++ b/docs/source/entity/per_world_mesh.rst
@ -0,0 +1,490 @@
+.. _heterogeneous_worlds:
+
+Heterogeneous Worlds
+====================
+
+mjlab can run a single batched simulation in which different parallel
+worlds use different mesh assets for the same logical entity. World 0
+may simulate a cube, world 1 a sphere, world 2 a bowl. All worlds
+share the same compiled scene and the same body and joint structure;
+only the meshes and the per-geom attributes that travel with them
+(friction, contact bits, mass, density, and a few more) differ across
+worlds. Articulated props work too (you can have a hinge or slide
+below the variant's root), as long as the joint topology matches
+across variants. The feature is exposed through ``VariantEntityCfg``.
+The full breakdown of what can and cannot vary across variants is in
+the next section.
+
+
+Quickstart
+----------
+
+Say you want some parallel worlds to hold a sphere and others to hold
+a cone, with a single shared scene running both at once. Define each
+variant as a function that returns an ``MjSpec``, then group them
+under one ``VariantEntityCfg``:
+
+.. code-block:: python
+
+    import mujoco
+
+    from mjlab.entity import EntityCfg, VariantEntityCfg
+
+
+    def make_sphere_spec() -> mujoco.MjSpec:
+        spec = mujoco.MjSpec()
+        mesh = spec.add_mesh(name="visual")
+        mesh.make_sphere(subdivision=3)
+        mesh.scale[:] = (0.05,) * 3
+        body = spec.worldbody.add_body(name="prop")
+        body.add_freejoint()
+        body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
+        return spec
+
+
+    def make_cone_spec() -> mujoco.MjSpec:
+        spec = mujoco.MjSpec()
+        mesh = spec.add_mesh(name="visual")
+        mesh.make_cone(nedge=16, radius=0.04)
+        body = spec.worldbody.add_body(name="prop")
+        body.add_freejoint()
+        body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
+        return spec
+
+
+    object_cfg = VariantEntityCfg(
+        variants={
+            "sphere": make_sphere_spec,
+            "cone":   make_cone_spec,
+        },
+        assignment={"cone": 2.0},  # twice as many cones as spheres
+        init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
+    )
+
+Plug the variant entity into a :ref:`scene` exactly like a regular
+``EntityCfg``:
+
+.. code-block:: python
+
+    from mjlab.scene import SceneCfg
+
+    scene_cfg = SceneCfg(
+        num_envs=4096,
+        entities={"object": object_cfg},
+    )
+
+Twice as many worlds will hold a cone as a sphere. Variants not listed
+in the ``assignment`` dict default to weight 1.0; omit ``assignment``
+entirely for uniform allocation across all variants.
+
+
+What variants can differ in
+---------------------------
+
+**Free to vary across variants:** the mesh asset assigned to each
+slot, the number of mesh geoms per ``(body, role)`` bucket on the
+variant body (one variant can have more collision meshes than
+another), the per-mesh-geom attributes that travel with the mesh
+(friction, contact bits, mass, density, ``condim``, and a handful of
+others), and explicit body inertial values within whichever single
+inertial mode the variants agree on per body.
+
+**Must match across variants:** the body tree, joint topology,
+primitive (non-mesh) geoms, and any actuators / sensors / tendons /
+equalities. Variants must also agree on the inertial representation
+per body (mesh-derived, diagonal, or fullinertia), and may not use the
+reserved ``mjlab/pad/`` name prefix on any element. Variant entities
+must also be floating-base: the root body declares a freejoint.
+
+The validator runs at entity build time and raises ``ValueError``
+naming the offending variant and the exact mismatch.
+
+
+How variants are assembled
+--------------------------
+
+mjlab merges every variant's mesh assets into a single ``MjSpec`` and
+gives the variant body enough mesh-geom *slots* to cover the maximum
+mesh count any variant uses for each ``(body, role)`` bucket. A slot
+is identified by ``(body_path, role, ordinal)``. ``role`` is "visual"
+or "collision", derived from ``contype``/``conaffinity``;
+mujoco_warp's ``geom_contype``/``geom_conaffinity`` are 1D shared
+(not per-world), so a slot's role is fixed across worlds by
+construction.
+
+A worked example
+~~~~~~~~~~~~~~~~
+
+Say variant ``sphere`` has 1 visual mesh geom and 2 collision mesh
+geoms on the prop body, and variant ``cone`` has 1 visual mesh geom
+and 4 collision mesh geoms on the same body.
+
+.. code-block:: text
+
+    sphere variant body              cone variant body
+    -------------------              -------------------
+    prop body                        prop body
+      [visual] sphere_vis              [visual] cone_vis
+      [coll]   sphere_col_0            [coll]   cone_col_0
+      [coll]   sphere_col_1            [coll]   cone_col_1
+                                       [coll]   cone_col_2
+                                       [coll]   cone_col_3
+
+mjlab walks each variant's body tree, buckets mesh geoms by
+``(body_path, role)``, and lays the union out as slots:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 8 18 8 12 27 27
+
+   * - Slot
+     - body_path
+     - role
+     - ordinal
+     - sphere fills with
+     - cone fills with
+   * - 0
+     - /prop
+     - visual
+     - 0
+     - sphere_vis
+     - cone_vis
+   * - 1
+     - /prop
+     - collision
+     - 0
+     - sphere_col_0
+     - cone_col_0
+   * - 2
+     - /prop
+     - collision
+     - 1
+     - sphere_col_1
+     - cone_col_1
+   * - 3
+     - /prop
+     - collision
+     - 2
+     - *(unfilled)*
+     - cone_col_2
+   * - 4
+     - /prop
+     - collision
+     - 3
+     - *(unfilled)*
+     - cone_col_3
+
+Five slots total. The merged scene's prop body has five mesh geoms:
+slot 0 plus four collision slots (the union of sphere's two and
+cone's four). At merge time, every variant's mesh asset is added to
+the merged spec under a unique name (e.g.
+``sphere/sphere_vis``, ``cone/cone_col_2``).
+
+The merged scene compiles once into a single canonical ``MjModel``
+that every world in the batch agrees on layout-wise: same nbody,
+ngeom, same body and geom IDs. mjlab's per-world overrides on top of
+that one model are what make worlds heterogeneous.
+
+What each world sees at runtime
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Worlds where ``sphere`` is active see only its three meshes; the two
+extra collision slots are disabled via per-world ``geom_dataid = -1``,
+and mujoco_warp skips them. Worlds where ``cone`` is active see all
+five meshes wired up.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 14 14 12 12 12 12 12
+
+   * - World
+     - variant
+     - slot 0
+     - slot 1
+     - slot 2
+     - slot 3
+     - slot 4
+   * - 0
+     - sphere
+     - sphere_vis
+     - sphere_col_0
+     - sphere_col_1
+     - **off (-1)**
+     - **off (-1)**
+   * - 1
+     - cone
+     - cone_vis
+     - cone_col_0
+     - cone_col_1
+     - cone_col_2
+     - cone_col_3
+
+Three categories of per-world override carry the variation:
+
+* **geom_dataid** is a ``(num_envs, ngeom)`` table. Its row for
+  world W picks which compiled mesh each slot points at. ``-1`` is
+  the "skip me" sentinel mujoco_warp already understands.
+* **Mesh-derived fields** (``geom_size``, ``geom_rbound``,
+  ``geom_aabb``, ``geom_pos``, ``geom_quat``, ``body_mass``,
+  ``body_subtreemass``, ``body_inertia``, ``body_invweight0``,
+  ``body_ipos``, ``body_iquat``) are stored as ``(num_envs, ...)``
+  arrays. The values for sphere worlds reflect a sphere-shaped
+  inertia tensor and sphere-sized AABBs; the values for cone worlds
+  reflect the cone. The full list is in
+  ``mjlab.entity.variants.VARIANT_DEPENDENT_FIELDS``.
+* **Per-mesh-geom attributes** (contact bits, friction, mass,
+  density, condim, group, priority, rgba, solref, solimp, margin,
+  gap) are captured per variant in ``VariantGeomSpec`` at merge time
+  and restored verbatim on the slot geom during the per-variant
+  reference compile. So if sphere's collision geoms have
+  ``friction=0.5`` and cone's have ``friction=1.2``, world W's
+  per-step friction reflects the assigned variant's source value.
+  The one exception is ``material``, which is not propagated across
+  variants; if you need per-world appearance variation use DR on
+  ``geom_rgba`` / ``mat_rgba``.
+
+If ``sphere`` adds a body that ``cone`` lacks (or vice versa), the
+validator rejects the configuration before any of the merge logic
+runs. The slot mechanism only flexes mesh geom counts within
+matching bodies; everything structural above the geom level must
+agree.
+
+.. note::
+
+   **Doesn't compiling the merged scene ruin the prop body's
+   inertia?**
+
+   No, but it's worth understanding why, because the naive intuition
+   says it should. If you stuck every variant's mesh geoms on the
+   prop body and called ``spec.compile()``, MuJoCo would sum each
+   geom's inertial contribution, and you would get a body whose mass
+   and inertia tensor are a meaningless mix of every variant's shape.
+
+   mjlab avoids this in two layers:
+
+   * **The merged scene does not stick every variant's geoms on the
+     body.** The prop body in the merged spec carries variant 0's
+     mesh geoms (with their original mass and density) plus, for any
+     slot variant 0 doesn't fill, a synthesized padding geom that has
+     ``mass = 0`` and ``density = 0``. Padding contributes nothing to
+     body inertia. Other variants' meshes are present in the merged
+     spec only as **mesh assets** (in the assets section, not as geoms
+     on any body). They get wired in at runtime via per-world
+     ``geom_dataid`` and never affect the host compile's inertial
+     sums.
+   * **Per-world overrides come from per-variant source compiles.**
+     Even with the above, the merged-scene compile's prop body inertia
+     is only correct for variant 0. For every other variant, mjlab
+     compiles that variant's original source spec in isolation (one
+     body, one variant's worth of meshes), reads the resulting
+     ``body_mass``, ``body_inertia``, ``body_ipos``, ``body_iquat``,
+     ``body_invweight0``, and ``body_subtreemass``, and writes them
+     into the per-world arrays at the prop body's index.
+
+   Net result: world W's prop body inertia is byte-equal to what you
+   would get by compiling variant W's source spec on its own. There
+   is a regression test
+   (``test_visual_collision_split_inertia_matches_independent_compile``
+   in ``tests/test_variants.py``) that asserts exactly this against
+   independent per-variant compiles.
+
+
+World assignment
+----------------
+
+How worlds get mapped to variants is controlled by the ``assignment``
+field on ``VariantEntityCfg``. It accepts three shapes:
+
+* ``None`` (default): uniform allocation across variants.
+* ``dict[str, float]``: per-variant weights. Variants not listed
+  default to weight 1.0.
+* ``Callable[[int], Sequence[int]]``: an explicit assignment function
+  called with ``num_envs`` at simulation init.
+
+Both the ``None`` and dict cases use the
+`largest remainder method
+<https://en.wikipedia.org/wiki/Largest_remainder_method>`_. Each
+variant's quota is ``q_i = (w_i / sum(w)) * num_envs``; each variant
+first receives ``floor(q_i)`` worlds, and the remaining
+``num_envs - sum(floors)`` worlds go to the variants with the largest
+fractional remainders, with ties broken by declaration order. For
+``num_envs = 10`` and weights ``(1.0, 2.0, 1.0)`` this gives
+``(3, 5, 2)`` worlds per variant. Weights are normalized internally,
+so ``{"a": 1, "b": 2, "c": 1}`` and ``{"a": 0.25, "b": 0.5, "c": 0.25}``
+produce identical assignments. A weight of zero is allowed and
+produces zero worlds for that variant; at least one variant must end
+up with positive weight.
+
+The default and dict paths are purely deterministic given
+``(assignment, num_envs)``. With ``assignment={"a": 1, "b": 1}`` and
+``num_envs = 8`` you always get ``[0, 0, 0, 0, 1, 1, 1, 1]``. There is
+no seed involved; rerunning the same config produces the same
+partition every time. Note that the partition's *boundaries* depend
+on ``num_envs``, so world W's variant is not necessarily stable when
+you change ``num_envs``. If you need explicit per-world stability
+across batch sizes (e.g. "world 0 is always variant 0, world 1 is
+always variant 1, regardless of how many envs I launch"), use a
+callable assignment as below.
+
+Variant assignment is fixed at ``Simulation`` initialization and does
+not resample on episode reset. The intended use is heterogeneous
+training across the batch, not per-episode mesh randomization.
+
+Read the resolved assignment from user code via
+``env.sim.world_to_variant``:
+
+.. code-block:: python
+
+    >>> env.sim.world_to_variant["object"]
+    tensor([0, 0, 0, 1, 1, 1, 1, 1, 1, 1])
+
+The mapping is keyed by entity name (without trailing slash) and
+returns a ``(num_envs,)`` tensor of variant indices in the order
+variants were declared in ``VariantEntityCfg.variants``. The dict is
+empty for non-variant scenes.
+
+
+Custom assignment with a callable
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When the weighted default is not what you want, pass a callable to
+``assignment``. The callable receives ``num_envs`` and must return a
+length-``num_envs`` sequence of variant indices in
+``[0, len(variants))``. The returned sequence's length and bounds are
+validated at sim init; mismatches raise a ``ValueError`` naming the
+offending entity.
+
+A few patterns:
+
+**Round-robin** - cycle through variants by world index.
+
+.. code-block:: python
+
+    cfg = VariantEntityCfg(
+        variants={"a": make_a, "b": make_b, "c": make_c},
+        assignment=lambda n: [w % 3 for w in range(n)],
+    )
+
+**Stratified halves** - first half is variant 0, second half is
+variant 1.
+
+.. code-block:: python
+
+    cfg = VariantEntityCfg(
+        variants={"easy": make_easy, "hard": make_hard},
+        assignment=lambda n: [0] * (n // 2) + [1] * (n - n // 2),
+    )
+
+Domain randomization
+--------------------
+
+Domain randomization on variant scenes preserves per-variant baselines
+automatically. When the simulation initializes, mjlab snapshots the
+variant-dependent fields as ``(num_envs, ...)`` tensors and registers
+them in ``sim.per_world_default_fields``. DR operations that read
+defaults (scale, additive offsets) detect this registration and index
+the per-world default array by environment, so a 10% mass scale
+applied across a batch containing a 100 g sphere variant and a 1 kg
+cube variant produces 10% perturbations *around each variant's own
+mass*, not 10% of a shared template mass.
+
+Fields that are not variant-dependent (``geom_friction``,
+``dof_armature``, ``dof_damping``, and so on) behave identically on
+variant and non-variant scenes.
+
+For inertial randomization the recommended path is
+``dr.pseudo_inertia``, which jointly randomizes mass, COM offset,
+principal moments of inertia, and principal frame orientation through
+the pseudo-inertia matrix factorization of `Rucker and Wensing (2022)
+<https://par.nsf.gov/servlets/purl/10347458>`_. It is exact for any
+perturbation magnitude and remains physically consistent across
+variants of different scale. ``dr.body_mass`` modifies ``body_mass``
+without touching the inertia tensor and emits a ``UserWarning`` when
+called; it is appropriate only for modeling a point mass added at the
+COM, not for density-like randomization. The distinction matters more
+on variant scenes than on single-asset scenes because variants often
+differ in mass by an order of magnitude.
+
+
+Viewers
+-------
+
+The native viewer, offscreen renderer, and Viser viewer all sync the
+selected environment's per-world fields into the host ``MjModel``
+before rendering, so the rendered geometry matches the variant
+assigned to the viewed environment. Switching environments in the
+native viewer (the ``,`` and ``.`` keys) updates the displayed mesh
+accordingly.
+
+Viser bakes mesh data into batched handles and cannot rely on a live
+view of ``geom_dataid``. It groups worlds by visual fingerprint (mesh
+selection, local geom frames, baked appearance) and builds one batched
+handle per group, with each environment assigned to its handle. A
+scene with N variants typically produces up to N handles per body.
+Convex hull visualization is computed per variant from the variant's
+mesh vertices.
+
+
+Performance
+-----------
+
+**Per-step cost is unaffected by variant count.** Variant-dependent
+fields are stored as per-world arrays accessed by world index in the
+existing kernels, with no branching or dispatch on variant.
+
+**Construction cost is linear in the total variant count.** mjlab
+compiles the merged scene once to produce the canonical ``MjModel``,
+then compiles each variant's original (un-merged) source spec in
+isolation to recover that variant's per-body and per-geom mesh-derived
+fields. Each per-variant compile sees only that variant's single body
+and mesh, so its cost is independent of the total number of variants
+in the scene.
+
+For a scene with one variant entity declaring k variants, construction
+runs ``1 + k`` compiles. With multiple variant entities, compiles
+decouple across entities: two variant entities of 5 variants each cost
+``1 + 5 + 5 = 11`` compiles, not ``1 + 5 * 5 = 26``. As an order of
+magnitude on CPU with typical procedural meshes, each per-variant
+compile takes around 1-2 ms, so a scene with 100 variants pays a few
+hundred milliseconds at startup and a scene with 1000 variants pays
+roughly two seconds.
+
+The merged spec contains every variant's mesh assets simultaneously,
+so memory at scene-build time scales with the total mesh vertex /
+face count across all variants. This is paid once at startup and does
+not affect training throughput.
+
+
+Limitations
+-----------
+
+**Floating-base only.** Each variant's root body must declare a free
+joint. Fixed-base variants are rejected; mocap auto-wrapping that
+applies to non-variant entities is not applied here.
+
+**Material assets are not propagated.** Each variant's ``contype``,
+``conaffinity``, ``condim``, ``friction``, ``mass``, ``density``,
+``group``, ``priority``, ``rgba``, ``solref``, ``solimp``, ``margin``,
+and ``gap`` are restored per-world during compile, but the
+``material`` reference on slot geoms inherits whichever material the
+template variant set. Use DR on ``geom_rgba`` / ``mat_rgba`` for
+per-world appearance variation.
+
+**Assignment is fixed at sim init.** There is no API to swap a world
+to a different variant on episode reset. World W's mesh asset is
+whatever it was assigned at init for the lifetime of the simulation.
+Per-episode mesh randomization is not supported today; DR can vary
+scalar properties (mass, friction, color, scale) on a fixed variant
+but cannot swap one mesh for another.
+
+**No support for per-world differing kinematic topology.** Variants
+must share the same body tree, joints, and actuator/sensor counts,
+so you cannot configure things like:
+
+* a different number of objects per world (world 0 has two props on
+  the table, world 1 has three);
+* different articulation per world (world 0's prop is an articulated
+  drawer with a slider joint, world 1's prop is a rigid block).
+
+True heterogeneous topology requires upstream support in mujoco_warp
+that does not currently exist.
--- a/docs/source/environment_config.rst
+++ b/docs/source/environment_config.rst
@ -0,0 +1,484 @@
+.. _environment_config:
+
+Environment Configuration
+=========================
+
+A single ``ManagerBasedRlEnvCfg`` dataclass fully specifies an mjlab
+environment: the physical world, the agent's interface to it, and the
+MDP defined on top. Because everything lives in one flat
+object, an environment can be inspected, copied, and modified without
+navigating a class hierarchy.
+
+For a broad orientation to mjlab before reading this page, start with
+:ref:`architecture_overview`.
+
+
+.. _env-config-skeleton:
+
+Annotated skeleton
+------------------
+
+The complete set of fields on ``ManagerBasedRlEnvCfg`` is shown below with
+inline comments. The fields marked with ``...`` must be provided; all others
+have defaults.
+
+.. code-block:: python
+
+    from dataclasses import dataclass, field
+
+    from mjlab.envs import ManagerBasedRlEnvCfg
+    from mjlab.managers.action_manager import ActionTermCfg
+    from mjlab.managers.command_manager import CommandTermCfg
+    from mjlab.managers.curriculum_manager import CurriculumTermCfg
+    from mjlab.managers.event_manager import EventTermCfg
+    from mjlab.managers.metrics_manager import MetricsTermCfg
+    from mjlab.managers.observation_manager import ObservationGroupCfg
+    from mjlab.managers.reward_manager import RewardTermCfg
+    from mjlab.managers.termination_manager import TerminationTermCfg
+    from mjlab.scene.scene import SceneCfg
+    from mjlab.sim.sim import SimulationCfg
+    from mjlab.viewer.viewer_config import ViewerConfig
+
+
+    @dataclass
+    class MyEnvCfg(ManagerBasedRlEnvCfg):
+
+        # --- Physics ---
+
+        decimation: int = 4
+        # Number of physics steps per policy step.
+        # Environment step duration = sim.mujoco.timestep * decimation.
+
+        sim: SimulationCfg = field(default_factory=SimulationCfg)
+        # Physics parameters: timestep, integrator, solver, contact settings.
+        # Default timestep is 0.002 s (500 Hz). Override with MujocoCfg.
+
+        scene: SceneCfg = ...
+        # Terrain, entities, and sensors. Also sets num_envs.
+        # Required; there is no default.
+
+        # --- Episode ---
+
+        episode_length_s: float = 20.0
+        # Episode duration in seconds.
+        # Steps = ceil(episode_length_s / (sim.mujoco.timestep * decimation)).
+
+        is_finite_horizon: bool = False
+        # False (default): time limit is an artificial cutoff. The agent
+        #   receives a truncated signal and bootstraps value beyond the limit.
+        # True: time limit defines the task boundary. The agent receives a
+        #   terminal done signal with no future value beyond it.
+
+        scale_rewards_by_dt: bool = True
+        # When True (default), each reward term is multiplied by step_dt so
+        # that cumulative episodic sums are invariant to simulation frequency.
+        # Set to False for algorithms that expect unscaled reward signals.
+
+        # --- Managers ---
+
+        observations: dict[str, ObservationGroupCfg] = field(default_factory=dict)
+        # Observation groups. Each key is a group name (e.g. "actor", "critic").
+        # Groups can differ in noise, history, delay, and concatenation.
+
+        actions: dict[str, ActionTermCfg] = field(default_factory=dict)
+        # Action terms. Each term controls one slice of the policy output
+        # and routes it to a specific entity's actuators.
+
+        rewards: dict[str, RewardTermCfg] = field(default_factory=dict)
+        # Reward terms. The manager computes a weighted sum each step.
+
+        terminations: dict[str, TerminationTermCfg] = field(default_factory=dict)
+        # Termination conditions. If empty, episodes never terminate early.
+        # Add a time_out term to enforce the episode length limit.
+
+        events: dict[str, EventTermCfg] = field(
+            default_factory=lambda: {
+                "reset_scene_to_default": EventTermCfg(
+                    func=reset_scene_to_default,
+                    mode="reset",
+                )
+            }
+        )
+        # Event terms for domain randomization and state resets.
+        # The default includes reset_scene_to_default, which resets all
+        # entities to their initial pose each episode. Override this dict
+        # to replace or extend the default reset behavior.
+
+        commands: dict[str, CommandTermCfg] = field(default_factory=dict)
+        # Command generators (e.g. velocity targets for locomotion).
+        # Commands are resampled at configurable intervals and on reset.
+
+        curriculum: dict[str, CurriculumTermCfg] = field(default_factory=dict)
+        # Curriculum terms that adjust training conditions based on performance.
+
+        metrics: dict[str, MetricsTermCfg] = field(default_factory=dict)
+        # Custom metrics logged as episode averages alongside reward terms.
+
+        # --- Misc ---
+
+        seed: int | None = None
+        # Random seed for reproducibility. If None, a random seed is chosen
+        # and stored back into this field after initialization.
+
+        viewer: ViewerConfig = field(default_factory=ViewerConfig)
+        # Camera position, resolution, and tracking target for rendering.
+
+
+.. _env-config-term-pattern:
+
+Term configuration pattern
+--------------------------
+
+All manager dictionaries follow the same pattern. Each entry maps a string
+name to a term configuration object. The configuration always carries at
+minimum a ``func`` field pointing to the callable that implements the term,
+and a ``params`` dict of extra keyword arguments forwarded to that callable.
+
+The manager calls ``func(env, **params)`` each step (or ``term(env, **params)``
+when ``func`` is a class that has been instantiated). Term names are arbitrary;
+they appear in training logs and are used only for identification.
+
+.. rubric:: Reward terms
+
+.. code-block:: python
+
+    from mjlab.envs import mdp
+    from mjlab.managers.reward_manager import RewardTermCfg
+    from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+    rewards = {
+        "alive": RewardTermCfg(
+            func=mdp.is_alive,
+            weight=1.0,
+        ),
+        "joint_torques": RewardTermCfg(
+            func=mdp.joint_torques_l2,
+            weight=-1e-4,
+            params={"asset_cfg": SceneEntityCfg("robot")},
+        ),
+        "action_rate": RewardTermCfg(
+            func=mdp.action_rate_l2,
+            weight=-0.1,
+        ),
+    }
+
+``weight`` scales the function's output before it is summed into the total
+reward. Negative weights produce penalties.
+
+``params`` maps to keyword arguments of the function. For example,
+``mdp.joint_torques_l2(env, asset_cfg=...)`` receives ``asset_cfg`` from the
+``params`` dict. Any argument not listed in ``params`` must have a default
+value in the function signature.
+
+.. rubric:: Termination terms
+
+.. code-block:: python
+
+    from mjlab.envs import mdp
+    from mjlab.managers.termination_manager import TerminationTermCfg
+
+    terminations = {
+        "time_out": TerminationTermCfg(
+            func=mdp.time_out,
+            time_out=True,   # marks this as a truncation, not a failure
+        ),
+        "fell_over": TerminationTermCfg(
+            func=mdp.bad_orientation,
+            params={"limit_angle": 1.22},   # ~70 degrees in radians
+        ),
+    }
+
+The ``time_out`` flag on ``TerminationTermCfg`` tells the manager to treat
+this condition as a truncation rather than a terminal failure. Truncations
+map to the ``truncated`` signal in the Gym interface; failures map to
+``terminated``. This distinction matters for value bootstrapping in RL
+algorithms.
+
+.. rubric:: Event terms
+
+.. code-block:: python
+
+    from mjlab.managers.event_manager import EventTermCfg
+
+    events = {
+        "reset_base": EventTermCfg(
+            func=mdp.reset_root_state_uniform,
+            mode="reset",
+            params={
+                "pose_range": {"yaw": (-3.14, 3.14)},
+                "velocity_range": {},
+            },
+        ),
+    }
+
+The ``mode`` field on ``EventTermCfg`` controls when the term fires:
+at startup, on episode reset, or at regular intervals. See :ref:`events`
+for the full treatment of lifecycle modes, built-in event functions, and
+the relationship between events and domain randomization.
+
+.. rubric:: Function-based vs. class-based terms
+
+Terms can be plain functions or classes. Functions are suitable for stateless
+computations; classes are useful when a term needs to cache expensive setup or
+maintain state across steps.
+
+A function-based term has the signature ``func(env, **params) -> Tensor``. A
+class-based term is instantiated once with ``(cfg, env)`` and then called with
+the same signature. Classes can optionally implement a ``reset(env_ids)`` hook
+for per-episode state clearing.
+
+.. code-block:: python
+
+    # Function-based (stateless)
+    RewardTermCfg(func=mdp.joint_torques_l2, weight=-0.01)
+
+    # Class-based (caches joint indices at init)
+    class MyReward:
+        def __init__(self, cfg, env):
+            self.joint_ids = resolve_joint_ids(cfg.params, env)
+
+        def __call__(self, env) -> torch.Tensor:
+            return compute_reward(env, self.joint_ids)
+
+    RewardTermCfg(func=MyReward, weight=1.0)
+
+
+.. _env-config-timing:
+
+Timing: decimation, timestep, and episode length
+-------------------------------------------------
+
+Three parameters jointly determine the temporal structure of the environment.
+
+``sim.mujoco.timestep``
+    The physics integration step in seconds. The default is 0.002 s (500 Hz).
+    This is one of the most important parameters in any environment: smaller
+    values produce more stable physics but slow down simulation. See the MuJoCo
+    `performance tuning <https://mujoco.readthedocs.io/en/stable/modeling.html#performance-tuning>`_
+    guide for practical advice on choosing timesteps and solver settings.
+
+``decimation``
+    The number of physics steps executed per policy step. The policy runs at
+    ``1 / (timestep * decimation)`` Hz.
+
+``episode_length_s``
+    The episode duration in seconds. The maximum number of policy steps per
+    episode is ``ceil(episode_length_s / (timestep * decimation))``.
+
+**Concrete example.** The velocity task uses ``timestep=0.005`` (200 Hz
+physics) and ``decimation=4``, giving a policy frequency of 50 Hz. With
+``episode_length_s=20.0``, each episode runs for exactly 1000 policy steps.
+
+.. code-block:: python
+
+    physics_dt  = 0.005        # seconds per physics step (200 Hz)
+    decimation  = 4            # physics steps per policy step
+    step_dt     = 0.005 * 4   # = 0.02 s per policy step (50 Hz)
+    episode_len = 20.0 / 0.02  # = 1000 policy steps per episode
+
+To read these values at runtime, use the environment properties:
+
+.. code-block:: python
+
+    env.physics_dt          # = cfg.sim.mujoco.timestep
+    env.step_dt             # = cfg.sim.mujoco.timestep * cfg.decimation
+    env.max_episode_length  # steps (int)
+    env.max_episode_length_s  # seconds (float)
+
+When ``scale_rewards_by_dt=True`` (the default), each reward term is
+multiplied by ``step_dt`` before being returned. A reward function that
+returns a constant value of 1.0 contributes ``step_dt`` per step and
+approximately ``episode_length_s`` over a full episode, regardless of how
+``decimation`` and ``timestep`` are set. Changing the simulation frequency
+without disabling this scaling leaves reward magnitudes unchanged.
+
+
+.. _env-config-subclassing:
+
+Subclassing pattern
+-------------------
+
+mjlab uses plain dataclass inheritance rather than deeply nested class
+hierarchies. To build a task-specific configuration, subclass
+``ManagerBasedRlEnvCfg`` and override fields.
+
+The recommended approach is to define the full configuration in a factory
+function, then call it from robot-specific configs that override only the
+fields that differ. The velocity task uses this pattern: ``make_velocity_env_cfg``
+returns a fully assembled ``ManagerBasedRlEnvCfg``, and each robot
+configuration calls the factory and patches in robot-specific values such
+as the scene, joint name patterns, and action scale.
+
+A condensed version of the factory illustrates the full assembly pattern:
+
+.. code-block:: python
+
+    import math
+    from dataclasses import replace
+
+    from mjlab.envs import ManagerBasedRlEnvCfg
+    from mjlab.envs.mdp import dr
+    from mjlab.envs.mdp.actions import JointPositionActionCfg
+    from mjlab.managers.event_manager import EventTermCfg
+    from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+    from mjlab.managers.reward_manager import RewardTermCfg
+    from mjlab.managers.scene_entity_config import SceneEntityCfg
+    from mjlab.managers.termination_manager import TerminationTermCfg
+    from mjlab.scene import SceneCfg
+    from mjlab.sim import MujocoCfg, SimulationCfg
+    from mjlab.tasks.velocity import mdp
+    from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
+    from mjlab.terrains import TerrainEntityCfg
+    from mjlab.terrains.config import ROUGH_TERRAINS_CFG
+    from mjlab.viewer import ViewerConfig
+
+
+    def make_velocity_env_cfg() -> ManagerBasedRlEnvCfg:
+
+        observations = {
+            "actor": ObservationGroupCfg(
+                terms={
+                    "base_lin_vel": ObservationTermCfg(
+                        func=mdp.builtin_sensor,
+                        params={"sensor_name": "robot/imu_lin_vel"},
+                    ),
+                    "joint_pos": ObservationTermCfg(func=mdp.joint_pos_rel),
+                    "command": ObservationTermCfg(
+                        func=mdp.generated_commands,
+                        params={"command_name": "twist"},
+                    ),
+                    # additional terms omitted for brevity
+                },
+                concatenate_terms=True,
+                enable_corruption=True,
+            ),
+            "critic": ObservationGroupCfg(
+                terms={...},
+                concatenate_terms=True,
+                enable_corruption=False,
+            ),
+        }
+
+        actions = {
+            "joint_pos": JointPositionActionCfg(
+                entity_name="robot",
+                actuator_names=(".*",),
+                scale=0.5,
+                use_default_offset=True,
+            )
+        }
+
+        commands = {
+            "twist": UniformVelocityCommandCfg(
+                entity_name="robot",
+                resampling_time_range=(3.0, 8.0),
+                ranges=UniformVelocityCommandCfg.Ranges(
+                    lin_vel_x=(-1.0, 1.0),
+                    lin_vel_y=(-1.0, 1.0),
+                    ang_vel_z=(-0.5, 0.5),
+                    heading=(-math.pi, math.pi),
+                ),
+            )
+        }
+
+        events = {
+            "reset_base": EventTermCfg(
+                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": {},
+                },
+            ),
+            "foot_friction": EventTermCfg(
+                mode="startup",
+                func=dr.geom_friction,
+                params={
+                    "asset_cfg": SceneEntityCfg("robot", geom_names=[]),
+                    "operation": "abs",
+                    "ranges": (0.3, 1.2),
+                },
+            ),
+            "push_robot": EventTermCfg(
+                func=mdp.push_by_setting_velocity,
+                mode="interval",
+                interval_range_s=(1.0, 3.0),
+                params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
+            ),
+        }
+
+        rewards = {
+            "track_linear_velocity": RewardTermCfg(
+                func=mdp.track_linear_velocity,
+                weight=2.0,
+                params={"command_name": "twist", "std": math.sqrt(0.25)},
+            ),
+            "dof_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-1.0),
+            "action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.1),
+        }
+
+        terminations = {
+            "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True),
+            "fell_over": TerminationTermCfg(
+                func=mdp.bad_orientation,
+                params={"limit_angle": math.radians(70.0)},
+            ),
+        }
+
+        return ManagerBasedRlEnvCfg(
+            decimation=4,
+            episode_length_s=20.0,
+            sim=SimulationCfg(
+                nconmax=35,
+                njmax=1500,
+                mujoco=MujocoCfg(timestep=0.005, iterations=10, ls_iterations=20),
+            ),
+            scene=SceneCfg(
+                terrain=TerrainEntityCfg(
+                    terrain_type="generator",
+                    terrain_generator=replace(ROUGH_TERRAINS_CFG),
+                    max_init_terrain_level=5,
+                ),
+                num_envs=1,
+            ),
+            observations=observations,
+            actions=actions,
+            commands=commands,
+            events=events,
+            rewards=rewards,
+            terminations=terminations,
+        )
+
+Robot-specific configs call this factory and patch fields using
+``dataclasses.replace`` or direct assignment. Common per-robot overrides
+include ``scene`` (to add the robot entity and sensors), joint name patterns
+inside ``SceneEntityCfg``, action ``scale``, and body names for reward terms.
+
+.. note::
+
+   Isaac Lab uses deeply nested ``__post_init__`` overrides for configuration
+   inheritance. mjlab avoids that pattern: each ``ManagerBasedRlEnvCfg`` is a
+   flat, inspectable dataclass. A misspelled field name raises a ``TypeError``
+   at construction rather than silently creating a new attribute. See
+   :ref:`migration_isaac_lab` for a full comparison.
+
+
+Where to go next
+----------------
+
+The remaining pages in the Manager Layer section cover each manager in
+detail:
+
+- :ref:`observations`: observation groups, the processing pipeline
+  (clip, scale, noise, delay, history), and built-in observation functions.
+- :ref:`actions`: action types and how the action manager routes policy
+  output to actuators.
+- :ref:`rewards`: reward terms and scaling by dt.
+- :ref:`terminations`: episode end conditions and the truncation/failure
+  distinction.
+- :ref:`commands`: command generators and goal-conditioned task setup.
+- :ref:`events`: the event manager lifecycle (startup, reset, interval).
+- :ref:`domain_randomization`: the full ``dr`` module for domain
+  randomization.
+- :ref:`curriculum`: difficulty progression based on policy performance.
+- :ref:`metrics`: custom per-step metrics logged as episode averages.
--- a/docs/source/events.rst
+++ b/docs/source/events.rst
@ -0,0 +1,212 @@
+.. _events:
+
+Events
+======
+
+The event manager executes hooks at specific points in the environment
+lifecycle. Any logic that should run at startup, on episode reset, or at
+regular intervals during training is registered as an event term. Common
+examples include resetting entities to an initial state, applying domain
+randomization to model parameters, pushing the robot with random velocity
+perturbations, and initializing robot state from a reference motion clip.
+All of these are configured through the same ``EventTermCfg`` interface,
+differing only in the ``mode`` field that controls when each term fires.
+
+Domain randomization, one of the most common uses of events, has its own
+dedicated reference page. See :ref:`domain_randomization` for the full
+``dr`` module, available functions, and internals.
+
+.. code-block:: python
+
+    from mjlab.envs.mdp import events as event_fns, dr
+    from mjlab.managers.event_manager import EventTermCfg
+    from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+    events = {
+        # Reset all entities to their default state each episode.
+        "reset_scene": EventTermCfg(
+            func=event_fns.reset_scene_to_default,
+            mode="reset",
+        ),
+        # Randomize foot friction once at startup.
+        "foot_friction": EventTermCfg(
+            func=dr.geom_friction,
+            mode="startup",
+            params={
+                "asset_cfg": SceneEntityCfg("robot", geom_names=[".*foot.*"]),
+                "ranges": (0.3, 1.2),
+                "operation": "abs",
+            },
+        ),
+        # Push the robot at random intervals during the episode.
+        "push_robot": EventTermCfg(
+            func=event_fns.push_by_setting_velocity,
+            mode="interval",
+            interval_range_s=(1.0, 3.0),
+            params={
+                "velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)},
+            },
+        ),
+        # Transient random impulses with duration and cooldown.
+        "impulse": EventTermCfg(
+            func=event_fns.apply_body_impulse,
+            mode="step",
+            params={
+                "force_range": (-50.0, 50.0),
+                "torque_range": (0.0, 0.0),
+                "duration_s": (0.1, 0.2),
+                "cooldown_s": (1.0, 3.0),
+                "asset_cfg": SceneEntityCfg("robot", body_names=("base",)),
+            },
+        ),
+    }
+
+
+Lifecycle modes
+---------------
+
+The ``mode`` field on ``EventTermCfg`` determines when the term fires. The
+four modes correspond to the timescales of an RL training run: once at
+process startup, once per episode, periodically within an episode, and on
+every environment step.
+
+``"startup"``
+    Fires once during environment initialization, after all managers are
+    constructed. Every environment receives the event simultaneously. This
+    mode is intended for parameters that should differ across environments
+    but remain fixed for the entire training run, such as link masses or
+    joint armatures randomized via the ``dr`` module.
+
+``"reset"``
+    Fires on every episode reset, for each environment being reset. This is
+    the most common mode. State initialization (writing the robot back to
+    its default pose) and episode-level domain randomization both belong
+    here.
+
+    The optional ``min_step_count_between_reset`` field prevents the term
+    from firing too frequently when episodes are very short. The term is
+    skipped for any environment that has not taken at least that many steps
+    since its last trigger. The first invocation always fires regardless.
+
+``"interval"``
+    Fires at regular time intervals during training, independent of episode
+    boundaries. The trigger frequency is controlled by ``interval_range_s``,
+    a ``(min, max)`` range in seconds. After each trigger the manager
+    samples a new wait time uniformly from that range. Each environment has
+    its own independent timer by default; setting ``is_global_time=True``
+    synchronizes all environments to a single shared timer. Interval events
+    are the natural home for mid-episode perturbations such as external
+    pushes or drifting model parameters.
+
+``"step"``
+    Fires on every environment step, for all environments. This mode is
+    intended for continuous effects that must be evaluated each step, such
+    as ``apply_body_impulse`` which manages its own internal duration and
+    cooldown timers. Because step events run every step, they should be
+    lightweight or manage their own activation logic internally to avoid
+    unnecessary computation.
+
+As with all manager terms, ``func`` points to the callable and ``params``
+holds keyword arguments forwarded to it alongside ``env`` and ``env_ids``.
+Any ``SceneEntityCfg`` values inside ``params`` are resolved once at
+manager construction (regex patterns are matched to model indices at that
+point, not on every call). Terms can be plain functions or classes; see
+:ref:`env-config-term-pattern` for the general pattern.
+
+
+Built-in event functions
+------------------------
+
+The functions below are available in ``mjlab.envs.mdp.events``.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Function
+     - Description
+   * - ``reset_scene_to_default``
+     - Resets all entities to their default states: root pose and velocity
+       for floating-base entities, mocap pose for fixed-base entities, and
+       joint positions and velocities for articulated entities. Environment
+       origins are applied automatically. This is the default event on
+       ``ManagerBasedRlEnvCfg``; most environments keep it and add
+       additional terms alongside it.
+   * - ``reset_root_state_uniform``
+     - Resets a single entity's root pose and velocity with uniform random
+       offsets from the default. Accepts ``pose_range`` and
+       ``velocity_range`` dictionaries with keys ``"x"``, ``"y"``,
+       ``"z"``, ``"roll"``, ``"pitch"``, ``"yaw"``. Orientation
+       perturbations compose with the default quaternion. For fixed-base
+       robots, this is the only way to position them at their environment
+       origins; without it they stack at the world origin.
+   * - ``reset_root_state_from_flat_patches``
+     - Places an entity on a randomly chosen flat terrain patch based on
+       the environment's assigned terrain level and type. Falls back to
+       ``reset_root_state_uniform`` when no flat patches are available.
+       Useful for locomotion tasks where robots should spawn on level
+       ground within their assigned sub-terrain.
+   * - ``reset_joints_by_offset``
+     - Resets joint positions and velocities by adding a uniform random
+       offset to the entity's defaults, clamped to soft joint limits.
+   * - ``push_by_setting_velocity``
+     - Adds a random velocity increment to the entity's current root
+       velocity, simulating an external push. Typically used with
+       ``mode="interval"`` to test disturbance rejection.
+   * - ``apply_external_force_torque``
+     - Applies random forces and torques to one or more bodies via the
+       MuJoCo external wrench mechanism.
+   * - ``apply_body_impulse``
+     - Applies transient external wrenches to bodies with configurable
+       duration and cooldown. Each environment independently samples a
+       random force direction and holds it for a sampled duration, then
+       waits through a cooldown before firing again. Supports an optional
+       ``body_point_offset`` to shift the application point away from the
+       center of mass. Includes built in debug visualization that draws
+       force arrows in the viewer. Use with ``mode="step"``.
+   * - ``randomize_terrain``
+     - Assigns each environment to a random sub-terrain row and column,
+       ignoring the curriculum. Useful for evaluation or play mode.
+
+
+Writing custom event terms
+--------------------------
+
+An event function takes ``env`` and ``env_ids`` as its first two arguments
+and any additional parameters from ``EventTermCfg.params``. It modifies
+simulation state in place and returns nothing. For terms that need
+expensive one-time setup (such as loading data from disk), use a class
+so that the setup runs once at construction rather than on every call.
+For example, the following custom event term resets the robot to a
+random pose sampled from a pre-recorded dataset:
+
+.. code-block:: python
+
+    import torch
+    from mjlab.managers.manager_base import ManagerTermBase
+    from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+    class ResetFromDataset(ManagerTermBase):
+        """Reset the robot to a random pose from a dataset."""
+
+        def __init__(self, cfg, env):
+            super().__init__(env)
+            self._robot = env.scene["robot"]
+            self._poses = torch.load(
+                cfg.params["dataset_path"],
+                map_location=env.device,
+            )
+
+        def __call__(self, env, env_ids, **kwargs):
+            # Sample with replacement: each env gets an independent pose.
+            indices = torch.randint(
+                len(self._poses), (len(env_ids),), device=env.device,
+            )
+            self._robot.write_joint_position_to_sim(
+                self._poses[indices], env_ids=env_ids,
+            )
+
+When a term needs to maintain state or perform expensive setup, implement
+it as a class. See :ref:`env-config-term-pattern` for the general
+pattern. For custom DR terms that write to model fields, see
+:ref:`domain_randomization`.
--- a/docs/source/faq.rst
+++ b/docs/source/faq.rst
@ -0,0 +1,546 @@
+.. _faq:
+
+FAQ & Troubleshooting
+=====================
+
+This page collects common questions about **platform support**, **performance**,
+**training stability**, and **visualization**, along with practical debugging
+tips and links to further resources.
+
+Platform Support
+----------------
+
+Does it work on macOS?
+~~~~~~~~~~~~~~~~~~~~~~
+
+Yes, but only with limited performance. mjlab runs on macOS
+using **CPU-only** execution through MuJoCo Warp.
+
+- **Training is not recommended on macOS**, as it lacks GPU acceleration.
+- **Evaluation works**, but is significantly slower than on Linux with CUDA.
+
+For serious training workloads, we strongly recommend **Linux with an NVIDIA GPU**.
+
+Does it work on Windows?
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+We have performed preliminary testing on **Windows** and **WSL**, but some
+workflows are not guaranteed to be stable.
+
+- Windows support may **lag behind** Linux.
+- Windows will be **tested less frequently**, since Linux is the primary
+  development and deployment platform.
+- Community contributions that improve Windows support are very welcome.
+
+CUDA Compatibility
+~~~~~~~~~~~~~~~~~~
+
+Not all CUDA versions are supported by MuJoCo Warp.
+
+- See `mujoco_warp#101 <https://github.com/google-deepmind/mujoco_warp/issues/101>`_
+  for details on CUDA compatibility.
+- **Recommended**: CUDA **12.4+** (for conditional execution support in CUDA
+  graphs).
+
+How do I run on CPU without touching the GPU?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Passing ``device="cpu"`` puts all mjlab computation on the CPU, but it does
+**not** stop Warp from initializing the GPU. The first time Warp's runtime
+comes up, it eagerly enumerates and creates a CUDA context on **every**
+visible device, regardless of which device you requested. So on a machine
+with a visible GPU, a ``device="cpu"`` run still claims VRAM.
+
+This happens inside Warp and cannot be prevented from Python once the
+package is imported. To keep the process entirely off the GPU, hide the
+devices from CUDA before launching:
+
+.. code-block:: bash
+
+   CUDA_VISIBLE_DEVICES="" uv run train.py ...
+
+With no visible CUDA devices, Warp initializes CPU-only and never allocates
+on the GPU. See `issue #949
+<https://github.com/mujocolab/mjlab/issues/949>`_ for background.
+
+Performance
+-----------
+
+Is it faster than Isaac Lab?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Based on our experience over the last few months, mjlab is **on par or
+faster** than Isaac Lab.
+
+What GPU do you recommend?
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- **RTX 40-series GPUs** (or newer)
+- **L40s, H100**
+
+Does mjlab support multi-GPU training?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Yes, mjlab supports **multi-GPU distributed training** using
+`torchrunx <https://github.com/apoorvkh/torchrunx>`_.
+
+- Use ``--gpu-ids "[0, 1]"`` (or ``--gpu-ids all``) when running the ``train``
+  command.
+- See the :doc:`training/distributed_training` for configuration details and examples.
+
+Training & Debugging
+--------------------
+
+My training crashes with NaN errors
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+A typical error when using ``rsl_rl`` looks like:
+
+.. code-block:: bash
+
+   RuntimeError: normal expects all elements of std >= 0.0
+
+This occurs when NaN/Inf values in the **physics state** propagate to the
+policy network, causing its output standard deviation to become negative or NaN.
+
+There are many possible causes, including potential bugs in **MuJoCo Warp**
+(which is still in beta). mjlab offers two complementary mechanisms to help
+you handle this:
+
+1. **For training stability** - NaN termination
+
+Add a ``nan_detection`` termination to reset environments that hit NaN:
+
+.. code-block:: python
+
+   from mjlab.envs.mdp import terminations as mdp_term
+   from mjlab.managers.termination_manager import TerminationTermCfg
+
+   # In your ManagerBasedRlEnvCfg subclass:
+   terminations = {
+      # Your other terminations...
+      "nan_term": TerminationTermCfg(func=mdp_term.nan_detection),
+   }
+
+This marks NaN environments as terminated so they can reset while training
+continues. Terminations are logged as
+``Episode_Termination/nan_term`` in your metrics.
+
+.. warning::
+
+   This is a **band-aid solution**. If NaNs correlate with your task objective
+   (for example, NaNs occur exactly when the agent tries to grasp an object),
+   the policy will never learn to complete that part of the task. Always
+   investigate the **root cause** using ``nan_guard`` in addition to this
+   termination.
+
+2. **For debugging** - NaN guard
+
+Enable ``nan_guard`` to capture the simulation state when NaNs occur:
+
+.. code-block:: bash
+
+   uv run train.py --enable-nan-guard True
+
+See the :doc:`NaN Guard documentation <debugging/nan_guard>` for details.
+
+The ``nan_guard`` tool makes it easier to:
+
+- Inspect the simulation state at the moment NaNs appear.
+- Build a minimal reproducible example (MRE).
+- Report potential framework bugs to the
+  `MuJoCo Warp team <https://github.com/google-deepmind/mujoco_warp/issues>`_.
+
+Reporting well-isolated issues helps improve the framework for everyone.
+
+How can I inspect the generated scene XML?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use the ``export-scene`` script to write the full scene (XML and mesh assets)
+to a directory:
+
+.. code-block:: bash
+
+    uv run export-scene g1 --output-dir /tmp/g1
+
+The exported ``scene.xml`` can be loaded directly in MuJoCo for visual
+inspection or diffing. This is useful for verifying that task configuration
+and physics are set up correctly, and for creating minimal reproducible
+examples to share with mjlab or MuJoCo Warp developers. The script accepts task IDs,
+entity aliases (``g1``, ``go1``, ``yam``), or arbitrary import paths. See
+:doc:`debugging/export_scene` for full details.
+
+My contact sensor misses collisions when using decimation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+With ``decimation > 1`` the physics runs multiple substeps per policy
+step. A brief contact (e.g. a self collision or an illegal ground touch)
+can appear and disappear within the substep loop, so by the time the
+sensor is read, ``found`` is zero and the event is invisible to
+rewards and terminations.
+
+Set ``history_length`` on the ``ContactSensorCfg`` equal to your
+decimation value. The sensor then stores force, torque, and distance
+for the last *N* substeps. Your reward or termination function can
+inspect the history to detect contacts that would otherwise be missed:
+
+.. code-block:: python
+
+    ContactSensorCfg(
+        name="self_collision",
+        ...,
+        fields=("found", "force"),
+        history_length=4,  # matches decimation=4
+    )
+
+    # In the reward/termination function:
+    force_mag = torch.norm(sensor.data.force_history, dim=-1)  # [B, N, H]
+    had_contact = (force_mag > 10.0).any(dim=1).any(dim=-1)    # [B]
+
+See :ref:`contact-sensor-history` for full details.
+
+.. note::
+
+   Feet ground sensors with ``track_air_time=True`` already accumulate
+   contact state across substeps, so they do not need history.
+
+.. _faq-sim-forward:
+
+When do I need to call ``sim.forward()``?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Short answer: you almost certainly don't.
+
+``sim.forward()`` wraps MuJoCo's ``mj_forward``, which runs the full forward
+dynamics pipeline (kinematics, contacts, forces, constraint solving, sensors)
+but skips integration, leaving ``qpos``/``qvel`` unchanged. It brings all
+derived quantities in ``mjData`` (``xpos``, ``xquat``, ``site_xpos``,
+``cvel``, ``sensordata``, etc.) into a consistent state with the current
+``qpos``/``qvel``.
+The environment's ``step()`` method calls it once per step, right before
+observation computation, so observations and commands always see fresh
+derived quantities. Termination, reward, and step/interval events run
+*before* this call and therefore see derived quantities that are stale by
+one physics substep, a deliberate tradeoff that avoids a second
+``forward()`` call while keeping the MDP well-defined (the staleness is
+consistent across all envs and all steps). Because events run before the
+call, any state they write (e.g. a velocity push) is refreshed by it and
+visible to the same step's observations.
+
+The one case where this matters is if you write an event or command that
+both writes state and reads derived quantities in the same function. For
+example, if Event A calls ``entity.write_root_velocity_to_sim()`` (which
+modifies ``qvel``) and then immediately reads ``entity.data.root_link_vel_w``
+(which comes from ``cvel``), the read will see stale values from before the
+write.
+
+.. warning::
+
+   Write methods (``write_root_state_to_sim``, ``write_joint_state_to_sim``,
+   etc.) modify ``qpos``/``qvel`` directly. Read properties
+   (``root_link_pose_w``, ``body_link_vel_w``, etc.) return derived
+   quantities that are only current as of the last ``sim.forward()`` call.
+   If you need to write then read in the same function, call
+   ``env.sim.forward()`` between them.
+
+For a deeper explanation, see `Discussion #289
+<https://github.com/mujocolab/mjlab/discussions/289>`_.
+
+Why aren't my training runs reproducible even with a fixed seed?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+MuJoCo Warp does not yet guarantee determinism, so running the same
+simulation with identical inputs may produce slightly different outputs.
+This is a known limitation being tracked in
+`mujoco_warp#562 <https://github.com/google-deepmind/mujoco_warp/issues/562>`_.
+
+Until determinism is implemented upstream, mjlab training runs will not be
+perfectly reproducible even when setting a seed.
+
+My XML ``<option>`` flags are not taking effect
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If you set simulation options like ``<flag contact="disable"/>`` in your
+entity XML, they will be silently ignored. This is because mjlab composes
+scenes by attaching entity specs into a parent scene spec using
+``MjSpec.attach()``, which does not propagate ``<option>`` settings from
+the child to the parent. This is a MuJoCo design decision: there is no
+sensible way to merge engine options (timestep, gravity, solver settings,
+etc.) across multiple attached models.
+
+To configure simulation options, use :class:`~mjlab.sim.sim.MujocoCfg` in
+your task's Python config:
+
+.. code-block:: python
+
+   from mjlab.sim.sim import MujocoCfg, SimulationCfg
+
+   sim=SimulationCfg(
+       mujoco=MujocoCfg(
+           disableflags=("contact",),
+           # timestep=0.01, gravity=(0, 0, -9.81), etc.
+       ),
+   )
+
+``MujocoCfg`` applies options directly to the compiled model, so they
+always take effect. mjlab will emit a warning if it detects non-default
+``<option>`` fields on an attached entity spec.
+
+Rendering & Visualization
+-------------------------
+
+What visualization options are available?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+mjlab currently supports two visualizers for policy evaluation and
+debugging:
+
+- **Native MuJoCo visualizer** - the built-in visualizer that ships with MuJoCo.
+- **Viser** - `Viser <https://github.com/nerfstudio-project/viser>`_,
+  a web-based 3D visualization tool.
+
+We are exploring **training-time visualization** (e.g., live rollout viewers),
+but this is not yet available.
+
+As an alternative, mjlab supports **video logging to Weights & Biases
+(W&B)**, so you can monitor rollout videos directly in the experiment dashboard.
+
+How many environments can I visualize at once?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Viewers render a small number of environments for performance reasons.
+
+- **Offscreen renderer** (for video recording): Renders the tracked
+  environment plus its nearest neighbors. The count is controlled by
+  ``ViewerConfig.max_extra_envs`` (default 2).
+- **Native/Viser viewers**: Limited by MuJoCo's geometry buffer
+  (default 10,000 geoms). The viewer shows whichever environments fit
+  within the geometry budget.
+
+Why are my fixed-base robots all stacked at the origin instead of in a grid?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Fixed-base robots require an **explicit reset event** to position them at
+their ``env_origins``. If your robots appear stacked at (0, 0, 0):
+
+**Common causes:**
+
+1. **Missing reset event** - Most common issue.
+2. **env_spacing is 0 or very small** - Check your ``SceneCfg(env_spacing=...)``.
+   Even with proper reset events, if ``env_spacing=0.0``, all robots will
+   be at the same position. If ``env_spacing`` is very small (e.g., 0.01),
+   they'll be clustered in a tiny area that looks like a line from a distance.
+
+**Solution**: Add a reset event that calls ``reset_root_state_uniform``:
+
+.. code-block:: python
+
+   # In your ManagerBasedRlEnvCfg
+   events = {
+     # For positioning the base of the robot at env_origins.
+     "reset_base": EventTermCfg(
+       func=mdp.reset_root_state_uniform,
+       mode="reset",
+       params={
+         "pose_range": {},  # Empty = use default pose + env_origins
+         "velocity_range": {},
+       },
+     ),
+     # ... other events
+   }
+
+This pattern is used in the example manipulation task (see ``lift_cube_env_cfg.py:85-94``).
+
+**Why this is needed**: Fixed-base robots are automatically wrapped in mocap
+bodies by ``auto_wrap_fixed_base_mocap()``, but mocap positioning only happens
+when you explicitly call a reset event. The ``env_origins`` offset is applied
+inside ``reset_root_state_uniform()`` at line 131 of ``envs/mdp/events.py``.
+
+See `issue #560 <https://github.com/mujocolab/mjlab/issues/560>`_ for examples.
+
+How does env_origins determine robot layout?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Robot spacing depends on your terrain configuration:
+
+**Plane terrain** (``terrain_type="plane"``):
+  - Creates an approximately square grid automatically
+  - Grid size: ``ceil(sqrt(num_envs))`` rows x cols
+  - Spacing controlled by ``env_spacing`` parameter (default: 2.0m)
+  - Examples with ``env_spacing=2.0``:
+    - 32 envs → 7x5 grid spanning 12m x 8m
+    - 4096 envs → 64x64 grid spanning 126m x 126m
+  - **Important**: If ``env_spacing=0``, all robots will be at (0, 0, 0)
+  - Implementation: ``terrain_importer.py:_compute_env_origins_grid()``
+
+**Procedural terrain** (``terrain_type="generator"``):
+  - Origins loaded from pre-generated terrain sub-patches
+  - Grid size: ``TerrainGeneratorCfg.num_rows x num_cols``
+  - Row index = difficulty level (curriculum mode)
+  - Column index = terrain type variant
+  - **Important allocation behavior**: Columns (terrain types) are evenly distributed
+    across environments, but rows (difficulty levels) are randomly sampled. This means
+    multiple environments can spawn on the same (row, col) patch, leaving others unoccupied,
+    even when ``num_envs > num_patches``.
+  - Example: 5x5 grid (25 patches), 100 envs → each column gets exactly 20 envs,
+    but those 20 are randomly distributed across 5 rows, so some patches remain empty.
+  - Supports ``randomize_env_origins()`` to shuffle positions during training
+
+How do I ensure each terrain type gets its own column?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Set ``curriculum=True`` in your ``TerrainGeneratorCfg``. This makes column
+allocation deterministic, with each column getting one terrain type based on
+normalized proportions.
+
+Example with 2 terrain types:
+
+.. code-block:: python
+
+   TerrainGeneratorCfg(
+     num_rows=3,
+     num_cols=2,
+     curriculum=True,  # Required for deterministic column allocation!
+     sub_terrains={
+       "flat": BoxFlatTerrainCfg(proportion=0.5),  # Gets column 0
+       "pillars": HfDiscreteObstaclesTerrainCfg(
+         proportion=0.5,  # Gets column 1
+       ),
+     },
+   )
+
+Without ``curriculum=True``, every patch is randomly sampled and you'll get
+a random mix of both terrain types scattered across all patches.
+
+**Note**: When ``num_cols`` equals the number of terrain types, each terrain
+gets exactly one column regardless of proportion values (they're normalized).
+When ``num_cols > num_terrain_types``, proportions determine how many columns
+each terrain type occupies.
+
+What is flat patch sampling and how does it affect robot spawning?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Flat patch sampling detects flat regions on heightfield terrains where robots
+can safely spawn. It uses morphological filtering on the heightfield to find
+circular areas where height variation is within a tolerance.
+
+Configure it on any sub-terrain via ``flat_patch_sampling``:
+
+.. code-block:: python
+
+   from mjlab.terrains.terrain_generator import FlatPatchSamplingCfg
+
+   "obstacles": HfDiscreteObstaclesTerrainCfg(
+     ...,
+     flat_patch_sampling={
+       "spawn": FlatPatchSamplingCfg(
+         num_patches=10,      # patches to sample per sub-terrain
+         patch_radius=0.5,    # flatness check radius (meters)
+         max_height_diff=0.05,  # max height variation within radius
+       ),
+     },
+   )
+
+Then use ``reset_root_state_from_flat_patches`` as your reset event to spawn
+robots on detected patches instead of at the sub-terrain center.
+
+**Key details:**
+
+- Only heightfield (``Hf*``) terrains support actual flat patch detection.
+  Box terrains (``Box*``) don't have heightfield data to analyze.
+- If any sub-terrain in the grid configures ``flat_patch_sampling``, the
+  flat patches array is allocated for **all** cells. Sub-terrains that don't
+  produce patches have their slots filled with the sub-terrain's spawn origin,
+  so ``reset_root_state_from_flat_patches`` always gets valid positions.
+- Without ``flat_patch_sampling``, use ``reset_root_state_uniform`` which
+  spawns at the sub-terrain origin (``env_origins``) plus an optional random
+  offset.
+
+Development & Extensions
+------------------------
+
+Can I develop custom tasks in my own repository?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Yes, mjlab has a **plugin system** that lets you develop tasks in separate
+repositories while still integrating seamlessly with the core:
+
+- Your tasks appear as regular entries for the ``train`` and ``play`` commands.
+- You can version and maintain your task repositories independently.
+
+A complete guide will be available in a future release.
+
+Assets & Compatibility
+----------------------
+
+What robots are included?
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+mjlab includes two **reference robots**:
+
+- **Unitree Go1** (quadruped).
+- **Unitree G1** (humanoid).
+
+These robots serve as:
+
+- Minimal examples for **robot integration**.
+- Stable, well-tested baselines for **benchmark tasks**.
+
+To keep the core library lean, we do **not** plan to aggressively expand the
+built-in robot library. Additional robots may be provided in separate
+repositories or community-maintained packages.
+
+Can I use USD or URDF models?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+No, mjlab expects **MJCF (MuJoCo XML)** models.
+
+- You will need to **convert** USD or URDF assets to MJCF.
+- For many common robots, you can directly use
+  `MuJoCo Menagerie <https://github.com/google-deepmind/mujoco_menagerie>`_,
+  which ships high-quality MJCF models and assets.
+
+Getting Help
+------------
+
+GitHub Issues
+~~~~~~~~~~~~~
+
+Use GitHub issues for:
+
+- **Bug reports**
+- **Performance regressions**
+- **Documentation gaps**
+
+When filing a bug, please include:
+
+- CUDA driver and runtime versions
+- GPU model
+- A minimal reproduction script
+- Complete error logs and stack traces
+- Appropriate labels (for example: ``bug``, ``performance``, ``docs``)
+
+`Open an issue <https://github.com/mujocolab/mjlab/issues>`_
+
+Discussions
+~~~~~~~~~~~
+
+Use GitHub Discussions for:
+
+- Usage questions (config, debugging, best practices)
+- Performance tuning tips
+- Asset conversion and modeling questions
+- Design discussions and roadmap ideas
+
+`Start a discussion <https://github.com/mujocolab/mjlab/discussions>`_
+
+Known Limitations
+-----------------
+
+We're tracking missing features for the stable release in
+https://github.com/mujocolab/mjlab/issues/100. Check our
+`open issues <https://github.com/mujocolab/mjlab/issues>`_ to see what's actively
+being worked on.
+
+If something isn't working or if we've missed something, please
+`file a bug report <https://github.com/mujocolab/mjlab/issues/new>`_.
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@ -0,0 +1,233 @@
+.. _installation:
+
+Installation Guide
+==================
+
+This guide presents different installation paths so you can
+choose the one that best fits your use case.
+
+.. contents::
+   :local:
+   :depth: 1
+
+.. note::
+
+    **System Requirements**
+
+    - **Training**: Linux + NVIDIA GPU (CUDA 12.4+ recommended)
+    - **Evaluation**: Linux, macOS, or Windows (WSL)
+    - **Python**: 3.10 or higher
+
+    See :ref:`faq` for more details on what is exactly supported.
+
+
+How to choose an installation method?
+-------------------------------------
+
+Select the card that best matches how you plan to use ``mjlab``.
+
+.. grid:: 2
+   :gutter: 2
+
+   .. grid-item-card:: Method 1 - Use mjlab as a dependency (uv)
+      :link: install-uv-dependency
+      :link-type: ref
+
+      You are **using mjlab as a dependency** in your own project managed by ``uv``. **(Recommended for most users)**
+
+   .. grid-item-card:: Method 2 - Develop / contribute (uv)
+      :link: install-uv-develop
+      :link-type: ref
+
+      You are **trying mjlab** or **contributing to mjlab itself** directly from inside the mjlab repository, with ``uv`` managing the environment.
+
+   .. grid-item-card:: Method 3 - Classic pip / venv / conda
+      :link: install-pip
+      :link-type: ref
+
+      You are using **classic tools** (``pip`` / ``venv`` / ``conda``) and **do not use uv**.
+
+   .. grid-item-card:: Method 4 - Docker / clusters
+      :link: install-docker
+      :link-type: ref
+
+      You are **running in containers or on clusters** and prefer a **Docker-based** setup.
+
+
+.. _install-uv-dependency:
+
+Method 1 - Use mjlab as a dependency (uv)
+-----------------------------------------
+
+This is our recommended way to use ``mjlab``. You have
+your own project and want to use ``mjlab`` as a dependency
+using ``uv``.
+
+1. Install uv
+^^^^^^^^^^^^^
+
+If you do not have ``uv`` installed, run:
+
+.. code-block:: bash
+
+   curl -LsSf https://astral.sh/uv/install.sh | sh
+
+2. Initialize your project
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Initialize a managed Python project:
+
+.. code-block:: bash
+
+   # Create a new package-based project
+   uv init --package my_mjlab_project
+   cd my_mjlab_project
+
+3. Add mjlab dependencies
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+There are different options to add ``mjlab`` as a dependency.
+We recommend using the latest stable version from PyPI. If you need
+the latest features, use the direct GitHub installation. Finally, if you
+need to use a feature you have developed locally, use the local editable
+install. These options are interchangeable: you can switch at any time.
+
+.. tab-set::
+
+   .. tab-item:: PyPI
+
+      Once in your project, install the latest snapshot from PyPI:
+
+      .. code:: bash
+
+         uv add mjlab
+
+   .. tab-item:: Source
+
+      Once in your project, install directly from GitHub without cloning:
+
+      .. code:: bash
+
+         uv add "mjlab @ git+https://github.com/mujocolab/mjlab"
+
+   .. tab-item:: Local
+
+      Clone the repository:
+
+      .. code:: bash
+
+         git clone https://github.com/mujocolab/mjlab.git
+
+      Once in your project, add it as an editable dependency:
+
+      .. code:: bash
+
+         uv add --editable /path/to/cloned/mjlab
+
+.. tip::
+
+   For a complete example of how to structure a project that integrates a custom robot
+   with an existing ``mjlab`` task, check out the
+   `ANYmal C Velocity Tracking <https://github.com/mujocolab/anymal_c_velocity>`_ repository.
+
+Verification
+^^^^^^^^^^^^
+
+After installation, verify that ``mjlab`` is working by running the demo:
+
+.. code-block:: bash
+
+   uv run demo
+
+
+.. _install-uv-develop:
+
+Method 2 - Develop / contribute (uv)
+------------------------------------
+
+This method is for developing ``mjlab`` itself or contributing to the project.
+
+.. code:: bash
+
+   git clone https://github.com/mujocolab/mjlab.git && cd mjlab
+   uv sync
+
+Verification
+^^^^^^^^^^^^
+
+After installation, verify that ``mjlab`` is working by running the demo:
+
+.. code-block:: bash
+
+   uv run demo
+
+
+.. _install-pip:
+
+Method 3 - Classic pip / venv / conda
+-------------------------------------
+
+Activate your virtual environment (``venv``, ``conda``, etc.), then install:
+
+.. code:: bash
+
+   pip install mjlab
+
+
+Verification
+^^^^^^^^^^^^
+
+After installation, verify that ``mjlab`` is working by running the demo:
+
+.. code-block:: bash
+
+   demo
+
+
+.. _install-docker:
+
+Method 4 - Docker / clusters
+----------------------------
+
+Prerequisites:
+
+- Install Docker: `Docker installation guide <https://docs.docker.com/engine/install/>`_.
+- Install an appropriate NVIDIA driver for your system and the
+  `NVIDIA Container Toolkit <https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html>`_.
+
+  - Be sure to register the container runtime with Docker and restart,
+    as described in the Docker configuration section of the NVIDIA
+    install guide.
+
+.. tab-set::
+
+   .. tab-item:: Pre-built image (recommended)
+
+      Pull and run the latest image from the GitHub Container Registry:
+
+      .. code-block:: bash
+
+         docker run --rm --runtime=nvidia --gpus all \
+           ghcr.io/mujocolab/mjlab uv run demo
+
+      The image is rebuilt on every push to ``main``.
+
+   .. tab-item:: Local build
+
+      Build from source and run:
+
+      .. code-block:: bash
+
+         ./scripts/run_docker.sh uv run demo
+
+
+Having some troubles?
+---------------------
+
+1. **Check the FAQ**
+
+    Consult the mjlab :ref:`faq` for answers to common installation and runtime issues
+
+2. **Still stuck?**
+
+    Open an issue on GitHub: https://github.com/mujocolab/mjlab/issues
--- a/Show More
+++ b/Show More