"""env.py — MJX environment for the modular legged robot. Multi-embodiment in one batched model: every env shares the master hexapod topology; the per-env MORPHOLOGY lives in model parameters (leg presence, link lengths, masses, actuator gains). Brax's randomization hook vmaps those parameters, so one PPO run trains on thousands of bodies at once. The policy sees a morphology vector derived from the same parameters, so it knows which body it is driving. Observation (93): gravity in body frame (3), gyro (3), command (3), per-leg morphology [present, cos, sin, thigh, shin] (6x5), joint pos (18), joint vel (18), previous action (18) Action (18): target joint angles, scaled; masked legs ignore theirs. """ from __future__ import annotations import jax, jax.numpy as jp import numpy as np import mujoco from mujoco import mjx from brax import base, envs from brax.envs.base import PipelineEnv, State from brax.io import mjcf as brax_mjcf import generator as G MAX_LEGS, JPL = G.MAX_LEGS, G.JOINTS_PER_LEG N_ACT = MAX_LEGS * JPL STAND = jp.tile(jp.array([0.0, 0.5, -1.0]), MAX_LEGS) # nominal pose per leg ACT_SCALE = jp.tile(jp.array([0.4, 0.7, 0.9]), MAX_LEGS) class ModBotEnv(PipelineEnv): def __init__(self, spec: G.BodySpec | None = None, cmd_range=(0.0, 0.35), yaw_range=0.6, episode_s=20.0, **kw): spec = spec or G.master_spec() mj_model = mujoco.MjModel.from_xml_string(G.to_mjcf(spec)) self.mj_model = mj_model sys = brax_mjcf.load_model(mj_model) n_frames = 5 # 0.004 s * 5 = 50 Hz super().__init__(sys=sys, backend="mjx", n_frames=n_frames, **kw) self.cmd_range, self.yaw_range = cmd_range, yaw_range self.max_steps = int(episode_s / (n_frames * 0.004)) name2id = lambda t, n: mujoco.mj_name2id(mj_model, t, n) self.torso_id = name2id(mujoco.mjtObj.mjOBJ_BODY, "torso") self.hip_ids = np.array([name2id(mujoco.mjtObj.mjOBJ_BODY, f"hip{i}") for i in range(MAX_LEGS)]) self.thigh_gid = np.array([name2id(mujoco.mjtObj.mjOBJ_GEOM, f"thigh{i}") for i in range(MAX_LEGS)]) self.shin_gid = np.array([name2id(mujoco.mjtObj.mjOBJ_GEOM, f"shin{i}") for i in range(MAX_LEGS)]) self.foot_gid = np.array([name2id(mujoco.mjtObj.mjOBJ_GEOM, f"foot{i}") for i in range(MAX_LEGS)]) self.floor_gid = name2id(mujoco.mjtObj.mjOBJ_GEOM, "floor") self.qpos0 = jp.array(mj_model.qpos0) self.lowers = jp.array(mj_model.jnt_range[1:, 0]) self.uppers = jp.array(mj_model.jnt_range[1:, 1]) # ---------- morphology, read from the (possibly vmapped) model ---------- def morphology(self, sys): present = (sys.actuator_gainprm[::JPL, 0] > 1e-3).astype(jp.float32) # abd actuator gain hip_pos = sys.body_pos[self.hip_ids] # (6,3) in torso frame ang = jp.arctan2(hip_pos[:, 1], hip_pos[:, 0]) thigh = 2.0 * sys.geom_size[self.thigh_gid, 1] / 0.1 # capsule half-length shin = 2.0 * sys.geom_size[self.shin_gid, 1] / 0.1 return jp.stack([present, jp.cos(ang), jp.sin(ang), thigh, shin], axis=1) # (6,5) # ---------- episode ---------- def sample_command(self, rng): r1, r2, r3 = jax.random.split(rng, 3) vx = jax.random.uniform(r1, (), minval=self.cmd_range[0], maxval=self.cmd_range[1]) vy = jax.random.uniform(r2, (), minval=-0.15, maxval=0.15) wz = jax.random.uniform(r3, (), minval=-self.yaw_range, maxval=self.yaw_range) return jp.array([vx, vy, wz]) def reset(self, rng: jax.Array) -> State: rng, r_cmd, r_q = jax.random.split(rng, 3) qpos = self.qpos0.at[7:].set(STAND + 0.05 * jax.random.normal(r_q, (N_ACT,))) qvel = jp.zeros(self.sys.nv) data = self.pipeline_init(qpos, qvel) info = {"rng": rng, "cmd": self.sample_command(r_cmd), "last_act": jp.zeros(N_ACT), "step": 0, "feet_air": jp.zeros(MAX_LEGS)} obs = self._obs(data, info) metrics = {k: 0.0 for k in ("r_lin", "r_ang", "p_vz", "p_rp", "p_torque", "p_rate", "p_pose", "alive", "fall")} return State(data, obs, jp.zeros(()), jp.zeros(()), metrics, info) def step(self, state: State, action: jax.Array) -> State: morph = self.morphology(self.sys) present_j = jp.repeat(morph[:, 0], JPL) target = STAND + ACT_SCALE * jp.clip(action, -1, 1) * present_j target = jp.clip(target, self.lowers, self.uppers) data = self.pipeline_step(state.pipeline_state, target) # body-frame quantities quat = data.xquat[self.torso_id] R = _quat_to_mat(quat) up = R[2, 2] gvec = R.T @ jp.array([0, 0, -1.0]) lin_w = data.cvel[self.torso_id, 3:] # world linear vel of torso lin_b = R.T @ lin_w ang_b = R.T @ data.cvel[self.torso_id, :3] cmd = state.info["cmd"] qpos = data.qpos[7:]; qvel = data.qvel[6:] torque = data.actuator_force # ---------- reward ---------- r_lin = jp.exp(-jp.sum((lin_b[:2] - cmd[:2]) ** 2) / 0.05) r_ang = jp.exp(-(ang_b[2] - cmd[2]) ** 2 / 0.25) p_vz = lin_b[2] ** 2 p_rp = jp.sum(ang_b[:2] ** 2) p_torque = jp.sum((torque * present_j) ** 2) p_rate = jp.sum(((action - state.info["last_act"]) * present_j) ** 2) p_pose = jp.sum(((qpos - STAND) * present_j) ** 2) fall = (up < 0.5) | (data.xpos[self.torso_id, 2] < 0.05) reward = (1.5 * r_lin + 0.5 * r_ang - 1.0 * p_vz - 0.05 * p_rp - 1e-4 * p_torque - 0.01 * p_rate - 0.05 * p_pose + 0.1) * 0.02 * 50 # per-second scale reward = jp.where(fall, -1.0, reward) step = state.info["step"] + 1 done = jp.where(fall | (step >= self.max_steps), 1.0, 0.0) # occasionally resample the command mid-episode rng, r_cmd, r_p = jax.random.split(state.info["rng"], 3) resample = jax.random.uniform(r_p) < 0.01 cmd = jp.where(resample, self.sample_command(r_cmd), cmd) info = dict(state.info, rng=rng, cmd=cmd, last_act=action, step=step) obs = self._obs(data, info) metrics = dict(state.metrics, r_lin=r_lin, r_ang=r_ang, p_vz=p_vz, p_rp=p_rp, p_torque=p_torque, p_rate=p_rate, p_pose=p_pose, alive=1.0 - done, fall=fall.astype(jp.float32)) return state.replace(pipeline_state=data, obs=obs, reward=reward, done=done, metrics=metrics, info=info) def _obs(self, data, info): morph = self.morphology(self.sys) R = _quat_to_mat(data.xquat[self.torso_id]) gvec = R.T @ jp.array([0, 0, -1.0]) ang_b = R.T @ data.cvel[self.torso_id, :3] present_j = jp.repeat(morph[:, 0], JPL) return jp.concatenate([ gvec, ang_b * 0.25, info["cmd"], morph.reshape(-1), (data.qpos[7:] - STAND) * present_j, data.qvel[6:] * 0.05 * present_j, info["last_act"] * present_j, ]) def _quat_to_mat(q): w, x, y, z = q return jp.array([ [1 - 2*(y*y + z*z), 2*(x*y - z*w), 2*(x*z + y*w)], [2*(x*y + z*w), 1 - 2*(x*x + z*z), 2*(y*z - x*w)], [2*(x*z - y*w), 2*(y*z + x*w), 1 - 2*(x*x + y*y)]]) # ---------- morphology randomization: one body per env ---------- def make_randomizer(env: ModBotEnv, mode="scaling", n_bodies=1000, seed=0): """Returns a brax randomization_fn. mode: 'single' -> every env is the default hexapod (specialist training) 'scaling' -> each env draws one of n_bodies fixed sampled bodies (the N in the scaling law) 'test:' -> every env is TEST_SET[name] """ rng = np.random.default_rng(seed) if mode == "single": pool = [G.hexapod()] elif mode.startswith("test:"): pool = [G.TEST_SET[mode.split(":")[1]]] else: pool = [G.sample_spec(rng) for _ in range(n_bodies)] pool_params = [_spec_to_params(env, s) for s in pool] stacked = jax.tree_util.tree_map(lambda *x: jp.stack(x), *pool_params) def randomization_fn(sys, rng): idx = jax.vmap(lambda k: jax.random.randint(k, (), 0, len(pool)))(rng) picked = jax.tree_util.tree_map(lambda x: x[idx], stacked) in_axes = jax.tree_util.tree_map(lambda x: None, sys) in_axes = in_axes.tree_replace({k: 0 for k in picked}) sys = sys.tree_replace(picked) return sys, in_axes return randomization_fn def _spec_to_params(env: ModBotEnv, spec: G.BodySpec): """Turn a BodySpec into the model-parameter overrides that realize it on the master model.""" m = env.mj_model body_mass = np.array(m.body_mass); body_inertia = np.array(m.body_inertia) body_pos = np.array(m.body_pos); geom_size = np.array(m.geom_size); geom_pos = np.array(m.geom_pos) site_pos = np.array(m.site_pos); gain = np.array(m.actuator_gainprm); bias = np.array(m.actuator_biasprm) nid = lambda t, n: mujoco.mj_name2id(m, t, n) B, GM, S = mujoco.mjtObj.mjOBJ_BODY, mujoco.mjtObj.mjOBJ_GEOM, mujoco.mjtObj.mjOBJ_SITE body_mass[env.torso_id] = spec.plate_mass for i, leg in enumerate(spec.legs): hip, thigh, shin = nid(B, f"hip{i}"), nid(B, f"thigh{i}"), nid(B, f"shin{i}") tg, sg, fg, fs = nid(GM, f"thigh{i}"), nid(GM, f"shin{i}"), nid(GM, f"foot{i}"), nid(S, f"foot{i}_site") a = np.radians(leg.port_angle_deg) if leg.present: body_pos[hip, :2] = leg.port_radius * np.array([np.cos(a), np.sin(a)]) body_pos[hip, 2] = 0.0 geom_size[tg, 1] = leg.thigh / 2; geom_pos[tg] = [leg.thigh / 2, 0, 0] body_pos[shin] = [leg.thigh, 0, 0] geom_size[sg, 1] = leg.shin / 2; geom_pos[sg] = [0, 0, -leg.shin / 2] geom_pos[fg] = [0, 0, -leg.shin]; site_pos[fs] = [0, 0, -leg.shin] geom_size[fg, 0] = 0.012 sc = leg.link_mass_scale for b, base_m in ((hip, 0.055), (thigh, 0.09), (shin, 0.085)): body_mass[b] = base_m * sc; body_inertia[b] = m.body_inertia[b] * sc for j in range(JPL): gain[i*JPL + j, 0] = G.SERVO_KP_DEFAULT; bias[i*JPL + j, 1] = -G.SERVO_KP_DEFAULT else: # masked: tiny, massless, lifted inside the plate, zero gain body_pos[hip] = [0, 0, 0.05] geom_size[tg, 1] = 0.0005; geom_pos[tg] = 0; body_pos[shin] = [0.001, 0, 0] geom_size[sg, 1] = 0.0005; geom_pos[sg] = 0; geom_pos[fg] = 0; site_pos[fs] = 0 geom_size[fg, 0] = 0.0005 for b in (hip, thigh, shin): body_mass[b] = 1e-3; body_inertia[b] = 1e-7 for j in range(JPL): gain[i*JPL + j, 0] = 0.0; bias[i*JPL + j, 1] = 0.0 return dict(body_mass=jp.array(body_mass), body_inertia=jp.array(body_inertia), body_pos=jp.array(body_pos), geom_size=jp.array(geom_size), geom_pos=jp.array(geom_pos), site_pos=jp.array(site_pos), actuator_gainprm=jp.array(gain), actuator_biasprm=jp.array(bias)) envs.register_environment("modbot", ModBotEnv)