"""Robot spec -> MJCF.
The spec JSON is the single source of truth: this module turns it into a MuJoCo
model, and bench/js/robot.js turns the same file into three.js meshes. If you
change the spec schema, change both.
"""
import json
import math
from pathlib import Path
def load_spec(path):
with open(path) as f:
spec = json.load(f)
spec["_path"] = str(path)
return spec
def _v(x):
return " ".join(f"{float(c):.6g}" for c in x)
def _capsule_fromto(g):
return f'fromto="{_v(g["from"])} {_v(g["to"])}"'
def _geom_xml(g, indent):
t = g["type"]
bits = [f'type="{t}"']
if t in ("capsule", "cylinder"):
bits.append(_capsule_fromto(g))
bits.append(f'size="{g["radius"]:.6g}"')
elif t == "box":
bits.append(f'size="{_v(g["size"])}"')
if "pos" in g:
bits.append(f'pos="{_v(g["pos"])}"')
elif t == "sphere":
bits.append(f'size="{g["radius"]:.6g}"')
if "pos" in g:
bits.append(f'pos="{_v(g["pos"])}"')
elif t == "ellipsoid":
# The right shape for a crowned tyre. A sphere gives the point contact
# we want, but it is as WIDE as it is tall - a 62.5mm wheel sphere is
# 125mm across and buries itself in whatever the wheel is bolted to. An
# ellipsoid is curved in both directions (so it still pivots) while
# staying the true width of the tyre.
bits.append(f'size="{_v(g["size"])}"')
if "pos" in g:
bits.append(f'pos="{_v(g["pos"])}"')
else:
raise ValueError(f"unsupported geom type {t!r}")
if "mass" in g:
bits.append(f'mass="{g["mass"]:.6g}"')
if "rgba" in g:
bits.append(f'rgba="{_v(g["rgba"])}"')
if "friction" in g:
bits.append(f'friction="{_v(g["friction"])}"')
if "name" in g:
bits.append(f'name="{g["name"]}"')
if g.get("collide") is False:
# visual-only geom: it still carries mass but never touches anything
bits.append('contype="0" conaffinity="0"')
return " " * indent + ""
def _site_xml(s, indent):
bits = [f'name="{s["name"]}"', f'pos="{_v(s.get("pos", [0, 0, 0]))}"']
bits.append(f'type="{s.get("type", "sphere")}"')
bits.append(f'size="{_v(s.get("size", [0.005]))}"')
bits.append('rgba="1 0.4 0 0.25"')
return " " * indent + ""
def _joint_xml(j, indent):
bits = [f'name="{j["name"]}"', f'type="{j.get("type", "hinge")}"']
if "axis" in j:
bits.append(f'axis="{_v(j["axis"])}"')
if j.get("continuous"):
bits.append('limited="false"')
elif "range_deg" in j:
lo, hi = j["range_deg"]
bits.append(f'range="{math.radians(lo):.6g} {math.radians(hi):.6g}"')
bits.append('limited="true"')
for key in ("damping", "armature", "frictionloss", "stiffness"):
if key in j:
bits.append(f'{key}="{j[key]:.6g}"')
return " " * indent + ""
def _camera_xml(c, indent):
# A body-mounted camera. Default orientation (xyaxes) looks forward along +x (the robot's
# front) with +z up — set xyaxes explicitly to aim it elsewhere. fovy in degrees (wide-angle).
bits = [f'name="{c["name"]}"', f'pos="{_v(c.get("pos", [0, 0, 0]))}"']
if "xyaxes" in c:
bits.append(f'xyaxes="{_v(c["xyaxes"])}"')
elif "euler" in c:
bits.append(f'euler="{_v(c["euler"])}"')
if "fovy" in c:
bits.append(f'fovy="{c["fovy"]:.6g}"')
return " " * indent + ""
def _body_xml(body, indent, is_root=False):
pad = " " * indent
out = [f'{pad}']
if is_root and body.get("float", False):
out.append(f'{pad} ')
if "joint" in body and body["joint"]:
out.append(_joint_xml(body["joint"], indent + 2))
for g in body.get("geoms", []):
out.append(_geom_xml(g, indent + 2))
for s in body.get("sites", []):
out.append(_site_xml(s, indent + 2))
for c in body.get("cameras", []):
out.append(_camera_xml(c, indent + 2))
for child in body.get("children", []):
out.append(_body_xml(child, indent + 2))
out.append(f"{pad}")
return "\n".join(out)
def motor_limits(spec, a):
"""Resolve an actuator's real limits from the motor part it names.
A BLDC's torque is Kt * current, and Kt is 8.27/Kv. Putting the motor in
the spec rather than a hand-written force_range means the model saturates
where the real hardware saturates, which is most of what decides whether a
learned policy transfers.
"""
lib = spec.get("motors", {})
m = lib.get(a.get("motor", ""))
if not m:
return None
kt = 8.27 / float(m["kv"])
ratio = float(a.get("gear_ratio", 1.0))
eff = float(m.get("efficiency", 0.9)) if ratio > 1 else 1.0
volts = float(m.get("voltage", 24.0))
# Prefer the manufacturer's own torque figures where they exist. Kt*I is a
# good estimate but it ignores saturation and their thermal rating, and the
# datasheet number is the one the part will actually be held to.
cont = m.get("rated_torque")
peak = m.get("max_torque")
if cont is None:
cont = kt * float(m["rated_current_a"])
if peak is None:
peak = kt * float(m.get("max_current_a", m.get("rated_current_a")))
return {
"kt": kt,
"ratio": ratio,
"tau_cont": cont * ratio * eff,
"tau_peak": peak * ratio * eff,
"no_load_rad_s": (float(m["kv"]) * volts * 2 * math.pi / 60.0) / ratio,
"motor": a.get("motor"),
"from_datasheet": m.get("rated_torque") is not None,
}
def _actuator_xml(a, spec, indent):
pad = " " * indent
mode = a.get("mode", "torque")
lim = motor_limits(spec, a)
bits = [f'name="{a["name"]}"', f'joint="{a["joint"]}"']
if mode == "torque":
tag = "motor"
bits.append(f'gear="{a.get("gear", 1.0):.6g}"')
cr = a.get("ctrl_range") or ([-lim["tau_cont"], lim["tau_cont"]] if lim
else [-1, 1])
bits.append(f'ctrlrange="{_v(cr)}"')
elif mode == "position":
tag = "position"
bits.append(f'kp="{a.get("kp", 30.0):.6g}"')
bits.append(f'kv="{a.get("kv", 1.0):.6g}"')
cr = a.get("ctrl_range") or _joint_range_rad(spec, a["joint"])
bits.append(f'ctrlrange="{_v(cr)}"')
elif mode == "velocity":
tag = "velocity"
bits.append(f'kv="{a.get("kv", 1.0):.6g}"')
bits.append(f'ctrlrange="{_v(a.get("ctrl_range", [-20, 20]))}"')
else:
raise ValueError(f"unsupported actuator mode {mode!r}")
# Position servos are bounded by the motor's PEAK torque, not its
# continuous rating. A body sprints on peak and survives on continuous:
# capping at the rated figure fenced off half the leg motor (tau_peak was
# computed above and never used) and made the sim reflexes weaker than
# the hardware's. Sustained load is a pricing problem - the reward's
# torque^2 term charges it - not a model ceiling. An explicit force_range
# in the spec (the wheels' driver limit) always wins.
if a.get("force_range"):
fr = a["force_range"]
elif lim:
key = "tau_peak" if mode == "position" else "tau_cont"
fr = [-lim[key], lim[key]]
else:
fr = None
if fr:
bits.append(f'forcerange="{_v(fr)}"')
return pad + f"<{tag} " + " ".join(bits) + "/>"
def _walk(body):
yield body
for c in body.get("children", []):
yield from _walk(c)
def find_joint(spec, name):
for b in _walk(spec["root"]):
j = b.get("joint")
if j and j["name"] == name:
return j
raise KeyError(f"no joint {name!r} in spec")
def _joint_range_rad(spec, name):
j = find_joint(spec, name)
if j.get("continuous") or "range_deg" not in j:
return [-10.0, 10.0]
lo, hi = j["range_deg"]
return [math.radians(lo), math.radians(hi)]
def _sensor_xml(s, indent):
pad = " " * indent
t = s["type"]
# IMU-style sensors attach to a site via site="...".
site_types = {"gyro", "accelerometer", "magnetometer", "touch",
"velocimeter", "rangefinder"}
# Frame sensors (position/orientation of a frame) use objtype/objname, NOT site.
frame_types = {"framepos", "framequat", "framezaxis", "framelinvel", "frameangvel"}
joint_types = {"jointpos", "jointvel", "jointactuatorfrc"}
if t in site_types:
ref = f'site="{s["site"]}"'
elif t in frame_types:
# Reference a site by default; allow an explicit objtype/objname override.
objtype = s.get("objtype", "site")
objname = s.get("objname", s.get("site"))
ref = f'objtype="{objtype}" objname="{objname}"'
elif t in joint_types:
ref = f'joint="{s["joint"]}"'
else:
raise ValueError(f"unsupported sensor type {t!r}")
return pad + f'<{t} name="{s["name"]}" {ref}/>'
def _prop_xml(p, indent):
"""Static scenery a task brings with it: steps, ramps, blocks."""
pad = " " * indent
bits = [f'name="{p["name"]}"', f'type="{p.get("type", "box")}"',
f'pos="{_v(p.get("pos", [0, 0, 0]))}"']
if p.get("type", "box") == "box":
bits.append(f'size="{_v(p["size"])}"')
else:
bits.append(f'size="{_v(p["size"])}"')
bits.append(f'rgba="{_v(p.get("rgba", [0.45, 0.47, 0.53, 1]))}"')
if "friction" in p:
bits.append(f'friction="{_v(p["friction"])}"')
return pad + ""
def spec_to_mjcf(spec, props=None):
"""world.floor may carry bank_deg / pitch_deg to tilt the ground.
It has to be baked in here rather than set later: MuJoCo precomputes the
frame of static world geoms, so writing geom_quat at runtime changes the
number but leaves the plane's normal pointing straight up, and nothing
happens. That silently produced identical results at 3, 6 and 10 degrees.
"""
sim = spec.get("sim", {})
world = spec.get("world", {})
floor = world.get("floor", {})
# THE PLAYGROUND HILLS: floor.hills swaps the plane for a heightfield
# (the undulation data is filled in after compile - MJCF carries only
# the shape; see RobotEnv). The geom keeps the name "floor" so every
# contact check in the fleet works unchanged.
if floor.get("hills"):
floor_asset = ('')
floor_geom = ('')
else:
floor_asset = ""
floor_geom = ('').format(
fs=floor.get("size", 12.0),
be=math.radians(float(floor.get("bank_deg", 0.0))),
pe=math.radians(float(floor.get("pitch_deg", 0.0))),
fr=_v(floor.get("friction", [1.0, 0.005, 0.0001])))
fsize = floor.get("size", 12.0)
bodies = _body_xml(spec["root"], 4, is_root=True)
scenery = "\n".join(_prop_xml(p, 4) for p in (props or []))
actuators = "\n".join(_actuator_xml(a, spec, 4) for a in spec.get("actuators", []))
sensors = "\n".join(_sensor_xml(s, 4) for s in spec.get("sensors", []))
# Ground-truth sensors used by the reward and the replay writer. These are
# deliberately separate from the robot's own sensor list: the policy never
# sees them.
root_name = spec["root"]["name"]
# objtype="xbody", not "body". In MuJoCo, "body" means the INERTIAL frame,
# whose axes are the principal axes of inertia and get reordered as the mass
# distribution changes. Widening the chassis silently rotated that frame, so
# "up" started pointing sideways and the robot read as fallen while standing
# perfectly still. "xbody" is the body frame we actually mean.
truth = "\n".join([
f' ',
f' ',
f' ',
f' ',
f' ',
])
return f"""
{floor_asset}
{floor_geom}
{scenery}
{bodies}
{actuators}
{sensors}
{truth}
"""
if __name__ == "__main__":
import sys
spec = load_spec(sys.argv[1] if len(sys.argv) > 1
else Path(__file__).parent.parent / "robots/wheeled_biped.json")
print(spec_to_mjcf(spec))