trainnr_mjlab

The platform's reinforcement-learning trainer: the walk tasks on mjlab, trained on the robot the project onboarded and identified, with the measured actuator physics and its intervals as the randomization basis. It is the layer trainnr's train_walk, evaluate_walk, export_deployment and related tools spawn; it imports trainnr and nothing imports it back.

Its own package and its own environment on purpose: mjlab pins torch, CUDA and the MuJoCo release it was built against, a third instrument beside the sim and train environments. The pin is exact (mjlab==1.6.0) and an upgrade is one deliberate change. Since 2026-09-22 pyproject.toml overrides mjlab's mujoco~=3.11 to mujoco==3.13.0 and mujoco-warp==3.13.0 (the Gaussian-splat ray tracer the scene loop renders with); every walk certificate names the instrument it was judged on.

Modules

Usage

The MCP tools build these command lines (trainnr/trainnr/mcp_actions.py); they run from a checkout the same way. On WSL2 prefix with --env-file ../trainnr/wsl.env.

cd trainnr-mjlab && uv sync --extra viz

# train: agent `smoke` is minutes and saves nothing, `g3` is the full recipe
# (the train_walk tool's recipe="smoke" and recipe="full")
uv run python -m trainnr_mjlab.walk_train --agent g3 --robot go2 \
    --project ../projects/<name> --iterations 1500 [--fit <fit@hash>] [--dr-span 0.10]

# evaluate a checkpoint: paired trials, exact intervals
uv run python -m trainnr_mjlab.walk_verdict ../projects/<name>/runs/<run>/model_1499.pt \
    --trials 40 --seed 1000 --robot go2 --project ../projects/<name>

# export for deployment: ONNX plus manifest
uv run python -m trainnr_mjlab.walk_export ../projects/<name>/runs/<run>/model_1499.pt \
    --project ../projects/<name> --robot go2 --name <deployment>

# the tests that run without a GPU
uv run python -m unittest discover -s tests -t .

Extras: bam installs Rhoban's own package as the parity reference for the kernel's tests (never a runtime dependency); viz installs the Rerun SDK the recorder streams to, pinned to the Studio's own release line.

Training needs a CUDA GPU, and so does reward_preview.py today (its device is cuda:0). Evaluation of a trained checkpoint runs on the CPU, slowly. The environment is about 6 GB; the first tool job that needs it (training, evaluation, export) prepares it.