Created
June 9, 2025 11:20
-
-
Save llandsmeer/f8159497fd739effe2341ad928e6fd56 to your computer and use it in GitHub Desktop.
Minimal JAX/MJX
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
<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