218 lines
8.1 KiB
Python
218 lines
8.1 KiB
Python
"""GPU-batched MuJoCo simulation using MJX (JAX backend).
|
|
|
|
MJX runs all environments in parallel on GPU via JAX, providing
|
|
~10-100x speedup over the CPU MuJoCoRunner for large env counts (1024+).
|
|
|
|
Requirements:
|
|
pip install 'jax[cuda12]' # NVIDIA GPU
|
|
pip install jax # CPU fallback
|
|
"""
|
|
|
|
import dataclasses
|
|
|
|
import structlog
|
|
import torch
|
|
|
|
try:
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import mujoco
|
|
from mujoco import mjx
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"MJX runner requires JAX and MuJoCo MJX. Install with:\n"
|
|
" pip install 'jax[cuda12]' # GPU\n"
|
|
" pip install jax # CPU\n"
|
|
) from e
|
|
|
|
import numpy as np
|
|
|
|
from src.core.env import BaseEnv
|
|
from src.core.runner import BaseRunner, BaseRunnerConfig
|
|
from src.runners.mujoco import MuJoCoRunner # reuse _load_model
|
|
|
|
log = structlog.get_logger()
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MJXRunnerConfig(BaseRunnerConfig):
|
|
num_envs: int = 1024
|
|
device: str = "cuda"
|
|
dt: float = 0.002
|
|
substeps: int = 20
|
|
action_ema_alpha: float = 0.2
|
|
|
|
|
|
class MJXRunner(BaseRunner[MJXRunnerConfig]):
|
|
"""GPU-batched MuJoCo runner using MJX (JAX).
|
|
|
|
Physics runs entirely on GPU via JAX; observations flow to
|
|
PyTorch through zero-copy DLPack transfers.
|
|
"""
|
|
|
|
def __init__(self, env: BaseEnv, config: MJXRunnerConfig):
|
|
super().__init__(env, config)
|
|
|
|
@property
|
|
def num_envs(self) -> int:
|
|
return self.config.num_envs
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
return torch.device(self.config.device)
|
|
|
|
# ── Initialization ───────────────────────────────────────────────
|
|
|
|
def _sim_initialize(self, config: MJXRunnerConfig) -> None:
|
|
# Step 1: Load CPU model (reuses URDF → MJCF → actuator injection)
|
|
self._mj_model = MuJoCoRunner._load_model(self.env.robot)
|
|
self._mj_model.opt.timestep = config.dt
|
|
self._nq = self._mj_model.nq
|
|
self._nv = self._mj_model.nv
|
|
self._nu = self._mj_model.nu
|
|
|
|
# Step 2: Put model on GPU
|
|
self._mjx_model = mjx.put_model(self._mj_model)
|
|
|
|
# Step 3: Default reset state on GPU
|
|
default_data = mujoco.MjData(self._mj_model)
|
|
default_qpos = self.env.get_default_qpos(self._nq)
|
|
if default_qpos is not None:
|
|
default_data.qpos[:] = default_qpos
|
|
mujoco.mj_forward(self._mj_model, default_data)
|
|
self._default_mjx_data = mjx.put_data(self._mj_model, default_data)
|
|
|
|
# Step 4: Initialise all environments with small perturbations
|
|
self._rng = jax.random.PRNGKey(42)
|
|
self._batch_data = self._make_batched_data(config.num_envs)
|
|
|
|
# Step 5: EMA ctrl state (on GPU as JAX array)
|
|
self._smooth_ctrl = jnp.zeros((config.num_envs, self._nu))
|
|
|
|
# Step 6: JIT-compile the hot-path functions
|
|
self._compile_jit_fns(config.substeps)
|
|
|
|
# Keep one CPU MjData for offscreen rendering
|
|
self._render_data = mujoco.MjData(self._mj_model)
|
|
|
|
log.info(
|
|
"mjx_runner_ready",
|
|
num_envs=config.num_envs,
|
|
substeps=config.substeps,
|
|
jax_devices=str(jax.devices()),
|
|
)
|
|
|
|
def _make_batched_data(self, n: int):
|
|
"""Create *n* environments with small random perturbations."""
|
|
self._rng, k1, k2 = jax.random.split(self._rng, 3)
|
|
pq = jax.random.uniform(k1, (n, self._nq), minval=-0.05, maxval=0.05)
|
|
pv = jax.random.uniform(k2, (n, self._nv), minval=-0.05, maxval=0.05)
|
|
|
|
default = self._default_mjx_data
|
|
model = self._mjx_model
|
|
|
|
def _init_one(pq_i, pv_i):
|
|
d = default.replace(
|
|
qpos=default.qpos + pq_i,
|
|
qvel=default.qvel + pv_i,
|
|
)
|
|
return mjx.forward(model, d)
|
|
|
|
return jax.vmap(_init_one)(pq, pv)
|
|
|
|
def _compile_jit_fns(self, substeps: int) -> None:
|
|
"""Pre-compile the two hot-path functions so the first call is fast."""
|
|
model = self._mjx_model
|
|
default = self._default_mjx_data
|
|
|
|
# ── Batched step (N substeps per call) ──────────────────────
|
|
@jax.jit
|
|
def step_fn(data):
|
|
def body(_, d):
|
|
return jax.vmap(mjx.step, in_axes=(None, 0))(model, d)
|
|
|
|
return jax.lax.fori_loop(0, substeps, body, data)
|
|
|
|
self._jit_step = step_fn
|
|
|
|
# ── Selective reset ─────────────────────────────────────────
|
|
@jax.jit
|
|
def reset_fn(data, smooth_ctrl, mask, rng):
|
|
rng, k1, k2 = jax.random.split(rng, 3)
|
|
ne = data.qpos.shape[0]
|
|
|
|
pq = jax.random.uniform(
|
|
k1, (ne, default.qpos.shape[0]), minval=-0.05, maxval=0.05,
|
|
)
|
|
pv = jax.random.uniform(
|
|
k2, (ne, default.qvel.shape[0]), minval=-0.05, maxval=0.05,
|
|
)
|
|
|
|
m = mask[:, None] # (num_envs, 1) broadcast helper
|
|
|
|
new_qpos = jnp.where(m, default.qpos + pq, data.qpos)
|
|
new_qvel = jnp.where(m, default.qvel + pv, data.qvel)
|
|
new_ctrl = jnp.where(m, 0.0, data.ctrl)
|
|
new_smooth = jnp.where(m, 0.0, smooth_ctrl)
|
|
|
|
new_data = data.replace(qpos=new_qpos, qvel=new_qvel, ctrl=new_ctrl)
|
|
new_data = jax.vmap(mjx.forward, in_axes=(None, 0))(model, new_data)
|
|
|
|
return new_data, new_smooth, rng
|
|
|
|
self._jit_reset = reset_fn
|
|
|
|
# ── Step / Reset ─────────────────────────────────────────────────
|
|
|
|
def _sim_step(self, actions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
# PyTorch → JAX (zero-copy on GPU via DLPack)
|
|
actions_jax = jnp.from_dlpack(actions.detach().contiguous())
|
|
|
|
# EMA smoothing (vectorised on GPU, no Python loop)
|
|
alpha = self.config.action_ema_alpha
|
|
self._smooth_ctrl = (
|
|
alpha * actions_jax + (1.0 - alpha) * self._smooth_ctrl
|
|
)
|
|
|
|
# Set ctrl & run N substeps for all environments
|
|
self._batch_data = self._batch_data.replace(ctrl=self._smooth_ctrl)
|
|
self._batch_data = self._jit_step(self._batch_data)
|
|
|
|
# JAX → PyTorch (zero-copy on GPU via DLPack, cast to float32)
|
|
qpos = torch.from_dlpack(self._batch_data.qpos.astype(jnp.float32))
|
|
qvel = torch.from_dlpack(self._batch_data.qvel.astype(jnp.float32))
|
|
return qpos, qvel
|
|
|
|
def _sim_reset(self, env_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
# Build boolean mask (fixed shape → no JIT recompilation)
|
|
mask = torch.zeros(
|
|
self.config.num_envs, dtype=torch.bool, device=self.device,
|
|
)
|
|
mask[env_ids] = True
|
|
mask_jax = jnp.from_dlpack(mask)
|
|
|
|
self._batch_data, self._smooth_ctrl, self._rng = self._jit_reset(
|
|
self._batch_data, self._smooth_ctrl, mask_jax, self._rng,
|
|
)
|
|
|
|
# Return only the reset environments' states
|
|
ids_np = env_ids.cpu().numpy()
|
|
rq = self._batch_data.qpos[ids_np].astype(jnp.float32)
|
|
rv = self._batch_data.qvel[ids_np].astype(jnp.float32)
|
|
return torch.from_dlpack(rq), torch.from_dlpack(rv)
|
|
|
|
# ── Rendering ────────────────────────────────────────────────────
|
|
|
|
def render(self, env_idx: int = 0) -> np.ndarray:
|
|
"""Offscreen render — copies one env's state from GPU to CPU."""
|
|
self._render_data.qpos[:] = np.asarray(self._batch_data.qpos[env_idx])
|
|
self._render_data.qvel[:] = np.asarray(self._batch_data.qvel[env_idx])
|
|
mujoco.mj_forward(self._mj_model, self._render_data)
|
|
|
|
if not hasattr(self, "_offscreen_renderer"):
|
|
self._offscreen_renderer = mujoco.Renderer(
|
|
self._mj_model, width=640, height=480,
|
|
)
|
|
self._offscreen_renderer.update_scene(self._render_data)
|
|
return self._offscreen_renderer.render()
|