api/robots.py

"""Robots that generated scenes can include, and the facts the model is given about them.

Generated scenes never redefine the robot. They `<include>` the Menagerie
model file, exactly like Menagerie's own scene.xml, so the robot is always the
real, validated one and the LLM only writes the environment around it.
"""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path

import mujoco


@dataclass(frozen=True)
class RobotSpec:
    name: str
    directory: str  # inside mujoco_menagerie
    model_file: str  # the file scenes include
    root_body: str  # must exist in a valid scene
    description: str
    kind: str = "humanoid"  # humanoid | quadruped | arm
    stand_height: float | None = 0.79  # root body height when standing (m); None for fixed-base arms
    fall_height: float | None = 0.35  # root below this counts as fallen; None: no fall detection
    home_qpos: tuple | None = None  # starting pose when the model file has no keyframe
    hold_kp: float = 0.0  # PD gains for "hold" on torque-controlled joints (0: derived from torque limits)
    hold_kd: float = 0.0
    spawn_note: str = "The robot spawns at the origin: keep other objects clear of a 0.5 m radius around (0, 0)"


ROBOTS: dict[str, RobotSpec] = {
    "unitree_g1": RobotSpec(
        name="unitree_g1",
        directory="unitree_g1",
        model_file="g1.xml",
        root_body="pelvis",
        description="Unitree G1 humanoid (29 DoF, rev 1.0) from MuJoCo Menagerie",
    ),
    "unitree_h1": RobotSpec(
        name="unitree_h1",
        directory="unitree_h1",
        model_file="h1.xml",
        root_body="pelvis",
        description="Unitree H1 humanoid (19 DoF, torque-controlled) from MuJoCo Menagerie",
        stand_height=0.96,
        fall_height=0.5,
        # h1.xml has no keyframe. Menagerie's scene.xml "home" pose (hip -0.4, knee 0.8, ankle -0.4) puts the
        # centre of mass over the heels, so it tips backwards under a joint-holding controller; shifting 0.15 rad
        # from hip to ankle (feet stay flat) moves it to mid-foot. With these gains it stands for 60 s.
        home_qpos=(0, 0, 0.97, 1, 0, 0, 0, 0, 0, -0.25, 0.8, -0.55, 0, 0, -0.25, 0.8, -0.55, 0, 0, 0, 0, 0, 0, 0, 0, 0),
        hold_kp=1000.0,
        hold_kd=20.0,
    ),
    "unitree_go2": RobotSpec(
        name="unitree_go2",
        directory="unitree_go2",
        model_file="go2.xml",
        root_body="base",
        description="Unitree Go2 quadruped (12 DoF, torque-controlled) from MuJoCo Menagerie",
        kind="quadruped",
        stand_height=0.27,
        fall_height=0.15,
    ),
    "franka_emika_panda": RobotSpec(
        name="franka_emika_panda",
        directory="franka_emika_panda",
        model_file="panda.xml",
        root_body="link0",
        description="Franka Emika Panda arm (7 DoF + gripper, fixed base) from MuJoCo Menagerie",
        kind="arm",
        stand_height=None,
        fall_height=None,
        spawn_note="The arm's base is fixed at the origin and reaches about 0.85 m: put objects to manipulate "
        "on a table or the floor within that reach (e.g. 0.4-0.6 m in front, +x), clear of the base",
    ),
}
ROBOT_NAMES = tuple(ROBOTS)


def robot_dir(menagerie_dir: Path, spec: RobotSpec) -> Path:
    return menagerie_dir / spec.directory


@lru_cache
def robot_facts(menagerie_dir: str, name: str) -> str:
    """Plain-text summary of the robot model, read from the model itself (not from the LLM's memory)."""
    spec = ROBOTS[name]
    path = robot_dir(Path(menagerie_dir), spec) / spec.model_file
    m = mujoco.MjModel.from_xml_path(str(path))

    def names(obj, n):  # unnamed elements (e.g. some robots' free joint) are skipped
        return [x for i in range(n) if (x := mujoco.mj_id2name(m, obj, i))]

    joints = []
    for i in range(m.njnt):
        name_i = mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_JOINT, i) or f"(unnamed joint {i})"
        if m.jnt_limited[i]:
            lo, hi = m.jnt_range[i]
            joints.append(f"{name_i} [{lo:.2f}, {hi:.2f}]")
        else:
            joints.append(name_i)
    keys = names(mujoco.mjtObj.mjOBJ_KEY, m.nkey)
    lines = [
        f"Robot: {spec.description}.",
        f'Include it with exactly: <include file="{spec.model_file}"/>  (the file sits next to the scene).',
        'It already defines: its bodies, joints, actuators, meshes (meshdir="assets"), '
        f"option integrator, and keyframe(s) {keys}. Do not redefine or rename any of them.",
        (
            f"Root body: {spec.root_body} (free joint), standing height about {spec.stand_height} m."
            if spec.stand_height
            else f"Root body: {spec.root_body} (fixed to the world at the origin)."
        )
        + f" Timestep {m.opt.timestep} s.",
        f"nq={m.nq}, nv={m.nv}, nu={m.nu} actuators.",
        "Bodies: " + ", ".join(names(mujoco.mjtObj.mjOBJ_BODY, m.nbody)[1:]),
        "Joints (radians): " + ", ".join(joints),
        "Actuators: " + ", ".join(names(mujoco.mjtObj.mjOBJ_ACTUATOR, m.nu)),
        "Sites: " + ", ".join(names(mujoco.mjtObj.mjOBJ_SITE, m.nsite)),
        f"{spec.spawn_note} unless the scenario says otherwise.",
    ]
    return "\n".join(lines)


def reset_to_robot_keyframe(model: mujoco.MjModel, data: mujoco.MjData, root_body: str | None, home_qpos=None) -> None:
    """Reset to the robot's first keyframe (or `home_qpos`) without disturbing the rest of the scene.

    Robot keyframes (e.g. the G1's "stand") only list the robot's own joints. When a scene adds
    free bodies, MuJoCo pads the keyframe with zeros for them, which would teleport every box
    or ball to the world origin, inside the robot. So the keyframe is applied to the joints of
    the robot's subtree only; every other joint keeps the scene's default pose (qpos0).
    """
    mujoco.mj_resetData(model, data)
    root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, root_body) if root_body else -1
    if model.nkey == 0 and (home_qpos is None or root < 0):
        return
    if root < 0:  # no robot: the scene's own keyframe describes the whole scene
        mujoco.mj_resetDataKeyframe(model, data, 0)
        return

    def in_robot(body: int) -> bool:
        while body > 0:
            if body == root:
                return True
            body = model.body_parentid[body]
        return False

    home = list(home_qpos) if model.nkey == 0 else None  # consumed joint by joint, in the robot's order
    for j in range(model.njnt):
        if not in_robot(model.jnt_bodyid[j]):
            continue
        start = model.jnt_qposadr[j]
        end = model.jnt_qposadr[j + 1] if j + 1 < model.njnt else model.nq
        if home is None:
            data.qpos[start:end] = model.key_qpos[0][start:end]
        elif len(home) >= end - start:
            data.qpos[start:end] = home[: end - start]
            del home[: end - start]
    if model.nu and model.nkey:
        data.ctrl[:] = model.key_ctrl[0]
    mujoco.mj_forward(model, data)