Skip to content

Instantly share code, notes, and snippets.

@llandsmeer
Created June 9, 2025 11:20
Show Gist options
  • Save llandsmeer/f8159497fd739effe2341ad928e6fd56 to your computer and use it in GitHub Desktop.
Save llandsmeer/f8159497fd739effe2341ad928e6fd56 to your computer and use it in GitHub Desktop.
Minimal JAX/MJX
import jax
import tqdm
from jax import numpy as jnp
from brax.io import mjcf
from brax.envs import PipelineEnv, State
from brax.io import html
from brax.generalized import pipeline
sys = mjcf.load('osc-mouse.xml')
action = None
step = jax.jit(lambda state, action: pipeline.step(sys, state, action))
state = pipeline.init(sys, sys.init_q,
jnp.zeros_like(sys.init_q))
states = [state]
for i in tqdm.tqdm(range(1000)):
t = i * 0.05
d = jnp.arange(12)/4
phi = jnp.sin(t+d)
state = step(state, 0.001*phi)
states.append(state)
with open('index.html', 'w') as f: print(html.render(sys, states, height=840), file=f)
<mujoco model="simple_mouse">
<compiler angle="degree" inertiafromgeom="true" settotalmass="14"/>
<option gravity="0 0 -9.81" timestep="0.01" iterations="4" />
<custom>
<numeric data="1000" name="constraint_limit_stiffness"/>
<numeric data="4000" name="constraint_stiffness"/>
<numeric data="10" name="constraint_ang_damping"/>
<numeric data="20" name="constraint_vel_damping"/>
<numeric data="0.5" name="joint_scale_pos"/>
<numeric data="0.2" name="joint_scale_ang"/>
<numeric data="0.0" name="ang_damping"/>
<numeric data="1" name="spring_mass_scale"/>
<numeric data="1" name="spring_inertia_scale"/>
<numeric data="15" name="solver_maxls"/>
</custom>
<asset>
<texture builtin="gradient" height="100" rgb1=".4 .5 .6" rgb2="0 0 0" type="skybox" width="100"/>
<!-- <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>-->
<texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
<texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
<material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
<material name="geom" texture="texgeom" texuniform="true"/>
</asset>
<default>
<joint armature=".000" damping="0.00" limited="true"/>
<geom contype="0" conaffinity="1" condim="3" density="0.0" friction="0 0.0 0" rgba="0.2 0.2 0.2 1.0"/>
<motor ctrllimited="true" ctrlrange="-1 1"/>
</default>
<size nstack="3000"/>
<worldbody>
<!-- Backbone -->
<body name="mouse" pos="0 0 0.05">
<geom name="spine" type="capsule" fromto="0 0 0 0.057 0 0" size="0.004" mass="1"/>
<geom name="head" type="capsule" fromto="0.063 0 0 0.085 0 0" size="0.012" mass="0"/>
<geom name="tail" type="capsule" fromto="0.0 0 0 -0.085 0 0" size="0.002" mass="0"/>
<!-- Left Hindlimb -->
<body pos="0 -0.01 0">
<joint name="lhip" type="hinge" axis="0 1 0" range="-50 30"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.022" size="0.002" mass="0.005"/>
<body pos="0 0 -0.022">
<joint name="lknee" type="hinge" axis="0 1 0" range="40 145"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.025" size="0.002" mass="0.005"/>
<body pos="0 0 -0.025">
<joint name="lankle" type="hinge" axis="0 1 0" range="-50 50"/>
<geom type="capsule" fromto="0 0 0 0.019 0 0" size="0.002" mass="0.001"/>
</body>
</body>
</body>
<!-- Right Hindlimb -->
<body pos="0 +0.01 0">
<joint name="rhip" type="hinge" axis="0 1 0" range="-50 30"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.022" size="0.002" mass="0.005"/>
<body pos="0 0 -0.022">
<joint name="rknee" type="hinge" axis="0 1 0" range="40 145"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.025" size="0.002" mass="0.005"/>
<body pos="0 0 -0.025">
<joint name="rankle" type="hinge" axis="0 1 0" range="-50 50"/>
<geom type="capsule" fromto="0 0 0 0.019 0 0" size="0.002" mass="0.001"/>
</body>
</body>
</body>
<!-- Left Forelimb -->
<body pos="0.057 -0.01 -0.006">
<joint name="lshoulder" type="hinge" axis="0 1 0" range="-14 40"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.013" size="0.002" mass="0.005"/>
<body pos="0 0 -0.013">
<joint name="lelbow" type="hinge" axis="0 1 0" range="-153 -60"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.017" size="0.002" mass="0.005"/>
<body pos="0 0 -0.017">
<joint name="lwrist" type="hinge" axis="0 1 0" range="-20 50"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.010" size="0.002" mass="0.001"/>
</body>
</body>
</body>
<!-- Right Forelimb -->
<body pos="0.057 +0.01 -0.006">
<joint name="rshoulder" type="hinge" axis="0 1 0" range="-14 40"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.013" size="0.002" mass="0.005"/>
<body pos="0 0 -0.013">
<joint name="relbow" type="hinge" axis="0 1 0" range="-153 -60"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.017" size="0.002" mass="0.005"/>
<body pos="0 0 -0.017">
<joint name="rwrist" type="hinge" axis="0 1 0" range="-50 50"/>
<geom type="capsule" fromto="0 0 0 0 0 -0.010" size="0.002" mass="0.001"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="lhip" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="lknee" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="lankle" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="rhip" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="rknee" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="rankle" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="lshoulder" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="lelbow" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="lwrist" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="rshoulder" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="relbow" gear="150"/>
<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="rwrist" gear="150"/>
</actuator>
</mujoco>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment