"""generator.py — module graph -> MJCF, for the modular legged robot. One function builds every body the hardware can build, plus continuous variation the hardware can't. Used in three places: training (sample many bodies), the digital twin (from the robot's self-description), and the fixed test set. Design choice that makes multi-embodiment GPU training possible: every model is the SAME "master" topology (a body plate with MAX_LEGS ports). A body with fewer legs is the master with some legs masked out: those legs get tiny geoms, negligible mass, locked joints and zero actuator gain. Same model, different parameters -> MJX can batch thousands of bodies in one step(). """ from __future__ import annotations import math from dataclasses import dataclass, field, asdict import numpy as np MAX_LEGS = 6 JOINTS_PER_LEG = 3 # --- physical constants matching the real platform (metres, kg) --- SERVO_MASS = 0.055 # STS3215 SERVO_KP_DEFAULT = 8.0 # N·m/rad position gain, retune after actuator ID SERVO_TORQUE_MAX = 2.9 # N·m (30 kg·cm stall) SERVO_VEL_MAX = 4.7 # rad/s (0.222 s / 60 deg) PLATE_RADIUS = 0.11 PLATE_THICK = 0.006 @dataclass class LegSpec: present: bool = True port_angle_deg: float = 0.0 # where on the plate rim it attaches port_radius: float = PLATE_RADIUS thigh: float = 0.080 # m shin: float = 0.090 # m (variants: 0.070, 0.110) link_mass_scale: float = 1.0 @dataclass class BodySpec: legs: list = field(default_factory=list) plate_mass: float = 0.9 # plate + board + wiring name: str = "modbot" @property def n_legs(self): return sum(l.present for l in self.legs) def to_dict(self): return asdict(self) # ---------------- canonical bodies ---------------- def _ports(n_ports=MAX_LEGS): """Default rim positions for the master plate. Ports 0..5 at 60° spacing, starting at +x (front) rotated so 0/1 are front-left/front-right etc.""" return [90 - 60 * i - 30 for i in range(n_ports)] # 60,0,-60,-120,-180,-240 PORT_ANGLES = _ports() def master_spec(): return BodySpec(legs=[LegSpec(True, a) for a in PORT_ANGLES], name="hexapod") def from_port_map(port_map: dict, shins: dict | None = None, name="twin"): """Build from the robot's self-description: {port_index: variant}. variant in {'short','default','long'} -> shin length.""" shin_len = {"short": 0.070, "default": 0.090, "long": 0.110} legs = [] for i, a in enumerate(PORT_ANGLES): if i in port_map: legs.append(LegSpec(True, a, shin=shin_len[port_map[i]])) else: legs.append(LegSpec(False, a)) return BodySpec(legs=legs, name=name) def quadruped(shin="default"): return from_port_map({0: shin, 2: shin, 3: shin, 5: shin}, name=f"quad_{shin}") def hexapod(shin="default"): return from_port_map({i: shin for i in range(6)}, name=f"hex_{shin}") TEST_SET = { "quad_default": quadruped("default"), "hex_default": hexapod("default"), "quad_long": quadruped("long"), "hex_short": hexapod("short"), "pentapod": from_port_map({0:"default",1:"default",2:"default",3:"default",5:"default"}, name="pentapod"), "quad_forward": from_port_map({0:"default",1:"default",2:"default",5:"default"}, name="quad_forward"), "hex_mixed": from_port_map({0:"long",1:"short",2:"default",3:"long",4:"short",5:"default"}, name="hex_mixed"), "tripod": from_port_map({1:"default",3:"default",5:"default"}, name="tripod"), "quad_missing": from_port_map({0:"default",2:"default",5:"default"}, name="quad_missing"), "hex_missing2": from_port_map({0:"default",1:"default",3:"default",4:"default"}, name="hex_missing2"), } def sample_spec(rng: np.random.Generator, min_legs=3, max_legs=6): """Random body from the module design space + continuous variation.""" n = int(rng.integers(min_legs, max_legs + 1)) present = set(rng.choice(MAX_LEGS, size=n, replace=False).tolist()) legs = [] for i, a in enumerate(PORT_ANGLES): if i in present: legs.append(LegSpec(True, a + rng.uniform(-8, 8), port_radius=PLATE_RADIUS * rng.uniform(0.9, 1.1), thigh=rng.uniform(0.065, 0.095), shin=rng.uniform(0.065, 0.115), link_mass_scale=rng.uniform(0.8, 1.3))) else: legs.append(LegSpec(False, a)) return BodySpec(legs=legs, plate_mass=float(rng.uniform(0.7, 1.3)), name=f"sample_{n}legs") # ---------------- MJCF ---------------- def to_mjcf(spec: BodySpec) -> str: """Emit MJCF. Masked legs are kept in the tree (same topology) but shrunk.""" L = [] L.append(f''' ''') for i, leg in enumerate(spec.legs): a = math.radians(leg.port_angle_deg) px, py = leg.port_radius * math.cos(a), leg.port_radius * math.sin(a) quat = f"{math.cos(a/2):.6f} 0 0 {math.sin(a/2):.6f}" on = leg.present th = leg.thigh if on else 0.001 sh = leg.shin if on else 0.001 mscale = leg.link_mass_scale if on else 0.01 sm = SERVO_MASS if on else 0.001 # masked legs: no contact, invisible, near-massless; same tree so MJX can batch bodies hide = '' if on else ' contype="0" conaffinity="0"' rgba_s = "0.2 0.2 0.25 1" if on else "0 0 0 0" rgba_l = "0.85 0.7 0.45 1" if on else "0 0 0 0" rgba_f = "0.9 0.5 0.3 1" if on else "0 0 0 0" servo = lambda pos, nm: (f'') L.append(f''' {servo("0.02 0 0", f"s{i}a")} {servo("0.02 0 0", f"s{i}b")} {servo("0 0 0", f"s{i}c")} ''') L.append(''' ''') for i, leg in enumerate(spec.legs): kp = SERVO_KP_DEFAULT if leg.present else 0.0 for j in ("abd", "hip", "knee"): L.append(f' ') L.append(''' ''') return "\n".join(L) def mask_vector(spec: BodySpec) -> np.ndarray: """1 for present legs, 0 for masked; the policy sees this.""" return np.array([1.0 if l.present else 0.0 for l in spec.legs], dtype=np.float32) def param_vector(spec: BodySpec) -> np.ndarray: """Per-leg morphology features fed to the policy: [present, cos, sin, thigh, shin].""" out = [] for l in spec.legs: a = math.radians(l.port_angle_deg) out += [float(l.present), math.cos(a), math.sin(a), l.thigh / 0.1, l.shin / 0.1] return np.array(out, dtype=np.float32) if __name__ == "__main__": import mujoco, sys for name, spec in TEST_SET.items(): m = mujoco.MjModel.from_xml_string(to_mjcf(spec)) d = mujoco.MjData(m) mujoco.mj_resetData(m, d) # stand pose for i, leg in enumerate(spec.legs): if leg.present: d.ctrl[3*i+1] = 0.5; d.ctrl[3*i+2] = -1.0 for _ in range(500): mujoco.mj_step(m, d) print(f"{name:14s} legs={spec.n_legs} mass={sum(m.body_mass):.2f} kg " f"z={d.qpos[2]:.3f} m up={d.xmat[1][8]:.3f}")