Go1 Joystick: A PyTorch-Native Quadruped Playground

A PyTorch reimplementation of MuJoCo Playground’s Unitree Go1 joystick simulation environment, powered by mujoco-torch.

Why reimplement MuJoCo Playground's Go1 task in PyTorch?

Robot learning requires collecting enormous numbers of simulated experiences. A policy chooses an action, the physics engine advances the robot, and this process is repeated across many environments and many training iterations. GPU-accelerated physics makes this practical by simulating many environments in parallel and keeping simulation data close to the policy and training code.

There are several common ways to do this, but each involves a trade-off:

• The MuJoCo C backend is the standard reference implementation, with a widely used model format, dynamics, and ecosystem. However, it is a native CPU simulator rather than a PyTorch tensor program. Using it with a GPU-based PyTorch policy requires a boundary between the simulator and the training loop, limiting batching, compilation, and gradients through the simulation.
• MJX provides a GPU-accelerated, batched implementation of MuJoCo physics. However, MJX is built for JAX. A PyTorch policy can interoperate with it, but the policy and simulator then use different tensor types, transformation systems, compilation mechanisms, and debugging tools.
• Other PyTorch-compatible GPU simulators avoid the framework boundary, but they use their own physics engines, model formats, and environment APIs. Moving to one can mean giving up compatibility with MuJoCo models and reference environments.

This leaves a missing combination: MuJoCo-compatible physics that runs as native PyTorch code.

mujoco-torch addresses this gap by reimplementing the MuJoCo/MJX physics path using PyTorch tensors and operations. The simulator can therefore participate in PyTorch batching, compilation, device placement, and automatic differentiation while retaining the MuJoCo model and state conventions.

This repository builds on mujoco-torch to reimplement MuJoCo Playground’s Unitree Go1 joystick environment. The result keeps the policy, environment state, and physics simulation inside one PyTorch workflow, while remaining aligned with a canonical MuJoCo robot-learning task. That makes it possible to use GPU-batched rollouts and torch.compile without moving the environment into JAX or abandoning MuJoCo compatibility.

Differentiable simulation is an additional benefit: when gradients through the physics are useful, losses defined on simulated behavior can be differentiated with respect to actions, controls, or model parameters. This supports workflows such as system identification, trajectory optimization, and model-based control.

Installation

Note

This release currently depends on a specific Git commit of mujoco-torch because the latest published version does not yet include the changes required by this project. The remote-install commands below use the v0.0.1 release tag. Once a new mujoco-torch version containing those changes is released, these instructions will be updated to use the normal package-manager install path.

Using uv

For a project managed with uv, install the current release from its Git tag:

uv add "go1-joystick-torch @ git+https://github.com/bsprenger/go1-joystick-torch.git@v0.0.1"

If you prefer to clone the repository locally, add the clone as a path dependency instead:

git clone https://github.com/bsprenger/go1-joystick-torch.git ../go1-joystick-torch
uv add "go1-joystick-torch @ ../go1-joystick-torch"

For active development against the local clone, use an editable dependency:

uv add --editable "../go1-joystick-torch"

Using pip

If you use pip, install the current release from its Git tag:

python -m pip install \
  "go1-joystick-torch @ git+https://github.com/bsprenger/go1-joystick-torch.git@v0.0.1"

Alternatively, clone the repository and install it from the local path:

git clone https://github.com/bsprenger/go1-joystick-torch.git ../go1-joystick-torch
python -m pip install ../go1-joystick-torch

For active development against the local clone, use an editable install:

python -m pip install --editable ../go1-joystick-torch

Development setup

For development within this repository, clone it first and run:

git clone https://github.com/bsprenger/go1-joystick-torch.git
cd go1-joystick-torch
make install

make install uses uv to create or synchronize the development environment and install the repository's pre-commit hooks.

Note

The current release only supports flat-terrain Go1 joystick locomotion. Terrain variation, observation noise, domain randomization, and native batched reset/step APIs will be included in an upcoming release.

Quick start

After installation, create an environment and step it with a PyTorch action:

import torch

from go1_joystick_torch import Joystick

env = Joystick()
observation, state = env.reset(seed=0)

action = torch.zeros(env.action_size, dtype=observation.dtype)
observation, state, reward, terminated, truncated = env.step(state, action)

For a full description of Joystick, JoystickState, and the reward API, see the API reference.

Development

The development workflow is for contributors and maintainers working from a source checkout. The Contributing guide covers the contribution process and local setup expectations.

The Makefile collects repeatable commands for environment setup, code quality checks, tests, verification, package builds, and documentation. Run make help to see all available targets.

make check          # formatting, linting, typing, dependency checks
make test           # fast unit and contract suite
make test-coverage  # test suite with coverage reporting
make render-go1     # open the Go1 viewer for visual mesh validation
make verify         # verify the environment against the canonical reference
make build          # build a wheel
make docs-test      # build documentation without warnings
make docs           # serve the documentation locally

The verification guide describes the developer-only pipeline for checking policy conversion, environment behavior, and rollout equivalence against the MuJoCo Playground reference. The diagnostics guide describes the developer-only tools for investigating numerical and backend differences. These tools are available from the repository checkout but are not included in the core package published to PyPI.