"""RobotEnv: turns a spec into a learnable task. Two rules kept strictly: 1. The policy only ever sees what the spec's sensors publish, at their real rate, with noise, bias and latency. No cheating on privileged state. 2. The reward may use ground truth, because the reward only exists at training time and never ships to the robot. """ import math import numpy as np import mujoco from mjcf import load_spec, spec_to_mjcf, find_joint from tasks import make_task # armed by the polish pass (learn_balance.measure_wave): accumulates # wave amplitude across every env stepped while set FEEL_ACCUM = None class SensorBank: """Zero-order-hold + noise + per-episode bias over MuJoCo sensordata.""" def __init__(self, model, spec, control_dt, rng): self.rng = rng self.control_dt = control_dt self.entries = [] for s in spec["sensors"]: if not s.get("in_obs", False): continue sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, s["name"]) if sid < 0: raise KeyError(f"sensor {s['name']!r} missing from compiled model") adr = model.sensor_adr[sid] dim = model.sensor_dim[sid] period = 1.0 / float(s.get("rate_hz", 1e9)) self.entries.append({ "name": s["name"], "slice": slice(adr, adr + dim), "dim": int(dim), "period": period, "noise": float(s.get("noise", 0.0)), "bias_mag": float(s.get("bias", 0.0)), "scale": float(s.get("scale", 1.0)), "bias": np.zeros(dim), "held": np.zeros(dim), "next_t": 0.0, }) self.dim = sum(e["dim"] for e in self.entries) def reset(self): for e in self.entries: e["bias"] = self.rng.uniform(-e["bias_mag"], e["bias_mag"], size=e["dim"]) e["held"] = np.zeros(e["dim"]) e["next_t"] = 0.0 def read(self, sensordata, t): out = np.empty(self.dim) i = 0 for e in self.entries: if t >= e["next_t"]: raw = np.asarray(sensordata[e["slice"]], dtype=float) noise = self.rng.normal(0.0, e["noise"], size=e["dim"]) if e["noise"] else 0.0 e["held"] = raw + e["bias"] + noise e["next_t"] = t + e["period"] out[i:i + e["dim"]] = e["held"] * e["scale"] i += e["dim"] return out def noisy(self, name, sensordata, raw): """The corrupted reading for a named sensor, matching what the policy sees, so the estimator never gets a cleaner signal than the network.""" for e in self.entries: if e["name"] == name: return e["held"] / (e["scale"] or 1.0) sl = raw.get(name) return np.asarray(sensordata[sl], dtype=float) if sl else np.zeros(3) def labels(self): out = [] for e in self.entries: if e["dim"] == 1: out.append(e["name"]) else: out += [f"{e['name']}.{c}" for c in "xyzw"[:e["dim"]]] return out class AttitudeEstimator: """Complementary filter over the robot's own noisy gyro and accelerometer. A single accelerometer sample cannot tell you which way is up while the robot is moving: it measures gravity minus the body's own acceleration, and a balancing machine is always accelerating. Measured on this robot, tilt inferred from one sample is out by 25 degrees at the median, on a machine that falls at 55. The gyro has the opposite problem: excellent short term, drifts long term. Blending them is what every real balancing robot does on its MCU, so doing it here is modelling the hardware rather than cheating. The estimator sees only the same noisy, biased signals the policy does. """ def __init__(self, alpha=0.98, dt=0.01, heading_alpha=0.995): self.alpha = alpha self.heading_alpha = heading_alpha self.dt = dt self.up = np.array([0.0, 0.0, 1.0]) self.heading = 0.0 # radians, from the magnetometer def reset(self): self.up = np.array([0.0, 0.0, 1.0]) self.heading = 0.0 def update(self, gyro, accel): # Predict: a world-fixed direction seen from a rotating body obeys # d(up)/dt = -omega x up. pred = self.up - np.cross(gyro, self.up) * self.dt n = np.linalg.norm(pred) pred = pred / n if n > 1e-9 else self.up # Correct toward the accelerometer, which is right on average but noisy # and corrupted whenever the robot accelerates. a = np.linalg.norm(accel) if a > 1e-6: meas = np.asarray(accel) / a # trust the accelerometer less when it disagrees with 1g trust = max(0.0, 1.0 - abs(a - 9.81) / 9.81) k = (1.0 - self.alpha) * trust blended = (1.0 - k) * pred + k * meas n = np.linalg.norm(blended) self.up = blended / n if n > 1e-9 else pred else: self.up = pred return self.up def update_heading(self, gyro, mag): """Yaw from the magnetometer, tilt-compensated, blended with the gyro. Gravity fixes roll and pitch but says nothing about yaw, so without a magnetometer heading drifts forever. The magnetometer is the only absolute yaw reference - and the least trustworthy sensor on the robot. """ # integrate the yaw rate about the estimated vertical self.heading += float(np.dot(gyro, self.up)) * self.dt if mag is None: return self.heading m = np.asarray(mag, dtype=float) n = np.linalg.norm(m) if n < 1e-9: return self.heading m = m / n # project the field into the horizontal plane defined by the estimated up horiz = m - self.up * np.dot(m, self.up) hn = np.linalg.norm(horiz) if hn < 1e-6: return self.heading horiz /= hn east = np.cross(self.up, np.array([1.0, 0.0, 0.0])) en = np.linalg.norm(east) if en < 1e-6: return self.heading east /= en north = np.cross(east, self.up) meas = math.atan2(float(np.dot(horiz, east)), float(np.dot(horiz, north))) err = (meas - self.heading + math.pi) % (2 * math.pi) - math.pi self.heading += (1.0 - self.heading_alpha) * err return self.heading class FusedIMU: """A BNO055/BNO085-class chip: fusion runs on the sensor, attitude comes out. Modelling the chip's OUTPUT rather than re-deriving it from raw gyro and accelerometer, because that is what the hardware actually provides. Honesty lives in the error model instead: white noise, a slow wander the chip cannot remove, a heading axis that is always worse than tilt, and extra heading error from six BLDC motors sitting near the sensor. """ def __init__(self, cfg, dt, rng): d = math.radians self.dt = dt self.rng = rng self.tilt_noise = d(float(cfg.get("tilt_noise_deg", 0.4))) self.tilt_drift = d(float(cfg.get("tilt_drift_deg", 1.0))) self.head_noise = d(float(cfg.get("heading_noise_deg", 1.5))) self.head_drift = d(float(cfg.get("heading_drift_deg", 3.0))) self.motor_coupling = d(float(cfg.get("motor_coupling_deg", 6.0))) self.latency = max(0, int(round(float(cfg.get("latency_ms", 10)) / 1000.0 / dt))) self.reset() def reset(self): # a fixed offset for this power-on, plus a slow random walk on top self.tilt_bias = self.rng.normal(0, self.tilt_drift, 3) self.head_bias = self.rng.normal(0, self.head_drift) self.walk = np.zeros(3) self.head_walk = 0.0 self.queue = [] def update(self, up_true, heading_true, drive): # slow wander: the error the chip cannot calibrate away self.walk = 0.999 * self.walk + self.rng.normal(0, self.tilt_drift * 0.02, 3) self.head_walk = 0.999 * self.head_walk + self.rng.normal(0, self.head_drift * 0.02) up = np.asarray(up_true, dtype=float) up = up + self.tilt_bias + self.walk + self.rng.normal(0, self.tilt_noise, 3) n = np.linalg.norm(up) up = up / n if n > 1e-9 else np.array([0.0, 0.0, 1.0]) head = (heading_true + self.head_bias + self.head_walk + self.rng.normal(0, self.head_noise) + self.motor_coupling * drive) self.queue.append((up, head)) if len(self.queue) > self.latency: return self.queue.pop(0) return self.queue[0] class RobotEnv: def __init__(self, spec_path, seed=0, randomise=None, cmd_scale=1.0, task=None, task_cfg=None, layout="union"): self.spec = load_spec(spec_path) self.rng = np.random.default_rng(seed) # The incentive is chosen first, because it may bring scenery with it # (a step to climb, a marker to drive at) that has to exist in the # model before it compiles. task_name = task or (self.spec.get("task", {}) or {}).get("type") self.task_obj = make_task(self.spec, task_name, task_cfg, layout=layout) self.layout = layout self.task_name = self.task_obj.name self.props = (self.task_obj.build_props() if hasattr(self.task_obj, "build_props") else []) self.xml = spec_to_mjcf(self.spec, props=self.props) self.model = mujoco.MjModel.from_xml_string(self.xml) self.data = mujoco.MjData(self.model) if self.model.nhfield: # the playground hills: gentle deterministic undulations - # crossing sine waves, normalized 0..1 (MuJoCo scales by the # hfield z size). Deterministic: every visit is the same # landscape. A cosine-blended pad at the centre stays flat # so the heap rests level at spawn. nr = int(self.model.hfield_nrow[0]) nc = int(self.model.hfield_ncol[0]) xs = np.linspace(0.0, 1.0, nc) ys = np.linspace(0.0, 1.0, nr) X, Y = np.meshgrid(xs, ys) h = (0.45 * np.sin(2 * np.pi * (2.1 * X + 0.4)) * np.cos(2 * np.pi * (1.7 * Y + 0.15)) + 0.3 * np.sin(2 * np.pi * (3.7 * X - 1.1 * Y)) + 0.25 * np.cos(2 * np.pi * (1.1 * X + 2.9 * Y + 0.3))) h = (h - h.min()) / max(1e-9, float(h.max() - h.min())) cx, cy = nc // 2, nr // 2 R = max(3.0, nc / 10.0) for r in range(nr): for c in range(nc): dd = math.hypot(c - cx, r - cy) / R if dd < 1.0: w = 0.5 * (1.0 - math.cos(math.pi * dd)) h[r, c] = h[r, c] * w + h[cy, cx] * (1.0 - w) self.model.hfield_data[:] = h.ravel() sim = self.spec["sim"] self.control_dt = 1.0 / float(sim.get("control_hz", 100)) self.n_sub = max(1, int(round(self.control_dt / self.model.opt.timestep))) # Body limits stay with the robot, not the incentive: how far this # machine can lean before it is falling over is a fact about the # machine, whatever you are paying it to do. limits = self.spec.get("limits", self.spec.get("task", {})) self.h_target = float(limits.get("target_height", 0.30)) self.min_h = float(limits.get("terminate", {}).get("min_height", 0.15)) self.min_up = math.cos(math.radians( float(limits.get("terminate", {}).get("max_tilt_deg", 55)))) self.max_steps = int(round(self.task_obj.episode_s / self.control_dt)) self.cmd_scale = cmd_scale self.task_obj.curriculum = cmd_scale dr = dict(self.spec.get("domain_randomisation", {})) if randomise is not None: dr["enabled"] = bool(randomise) self.dr = dr self.ctrl_lo = self.model.actuator_ctrlrange[:, 0].copy() self.ctrl_hi = self.model.actuator_ctrlrange[:, 1].copy() # For a position servo ctrlrange is an angle, so the torque ceiling has # to come from forcerange. Reading ctrlrange gives you 1.571 "Nm" for a # +-90 degree hip, which is nonsense. self.tau_limit = np.abs(self.model.actuator_forcerange[:, 1]).copy() self.tau_limit[self.tau_limit < 1e-9] = 1.0 self.n_act = self.model.nu # Leg joints get a nominal-posture penalty; wheels are free to spin. self.leg_joint_ids, self.q_nominal = [], [] init_q = dict(self.spec.get("init", {}).get("qpos", {})) if getattr(self.task_obj, "pose", None): init_q.update(self.task_obj.pose) # this incentive stands differently self.init_q = init_q for name, val in init_q.items(): jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) self.leg_joint_ids.append(self.model.jnt_qposadr[jid]) self.q_nominal.append(float(val)) self.leg_joint_ids = np.array(self.leg_joint_ids, dtype=int) self.q_nominal = np.array(self.q_nominal) self.leg_dof_ids = np.array([ self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, n)] for n in init_q ], dtype=int) # Action mapping. A position servo on a joint with a nominal angle is # centred on that angle, so action 0 means "hold the standing pose" # rather than "middle of the mechanical travel". Everything else spans # its ctrlrange. self.act_center = 0.5 * (self.ctrl_lo + self.ctrl_hi) self.act_span = 0.5 * (self.ctrl_hi - self.ctrl_lo) # A mirrored motor turns the opposite way for the same command. self.act_sign = np.array([-1.0 if a.get("reversed") else 1.0 for a in self.spec.get("actuators", [])]) for i, a in enumerate(self.spec.get("actuators", [])): if a.get("mode") != "position": continue nominal = init_q.get(a["joint"]) if nominal is None: continue span = float(a.get("action_scale", 0.6)) self.act_center[i] = nominal self.act_span[i] = span # Settle the model in THIS task's pose to find where it actually stands. # Must come after init_q and the action mapping exist, or it settles a # different robot than the one the task will run. self._settle_height() # An incentive may expose only some actuators. The rest hold their # nominal pose, and the policy never sees or searches them. names = [a["name"] for a in self.spec.get("actuators", [])] want = getattr(self.task_obj, "actuators", None) self.act_idx = ([names.index(n) for n in want] if want else list(range(self.n_act))) self.n_act = len(self.act_idx) self._resolve_motors() self.sensors = SensorBank(self.model, self.spec, self.control_dt, self.rng) # Onboard attitude estimate, fed the same noisy sensors the policy gets. est_cfg = self.spec.get("estimator", {"enabled": True, "alpha": 0.98}) self.est_on = bool(est_cfg.get("enabled", True)) self.est_mode = est_cfg.get("type", "complementary") if self.est_mode == "fused": self.estimator = FusedIMU(est_cfg, self.control_dt, self.rng) else: self.estimator = AttitudeEstimator(float(est_cfg.get("alpha", 0.98)), self.control_dt, float(est_cfg.get("heading_alpha", 0.995))) self.has_mag = any(x.get("name") == "mag" and x.get("in_obs") for x in self.spec.get("sensors", [])) self.mag_coupling = next((float(x.get("motor_coupling", 0.0)) for x in self.spec.get("sensors", []) if x.get("name") == "mag"), 0.0) # up vector, plus heading as sin/cos when a magnetometer is fitted self.est_dim = 5 if self.est_on else 0 # up vector + heading sin/cos self._raw = {} for nm in ("gyro", "accel", "mag"): sid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SENSOR, nm) if sid >= 0: self._raw[nm] = slice(self.model.sensor_adr[sid], self.model.sensor_adr[sid] + self.model.sensor_dim[sid]) self.obs_dim = (self.sensors.dim + self.est_dim + self.task_obj.obs_extra + self.n_act) self.act_dim = self.n_act self.act_dim = self.n_act # Pristine copies for domain randomisation to perturb from. self._base_mass = self.model.body_mass.copy() self._base_inertia = self.model.body_inertia.copy() self._base_ipos = self.model.body_ipos.copy() self._base_friction = self.model.geom_friction.copy() self._base_gain = self.model.actuator_gainprm.copy() self._base_bias = self.model.actuator_biasprm.copy() self._base_forcerange = self.model.actuator_forcerange.copy() self._gt = {} for nm in ("gt_quat", "gt_pos", "gt_up", "gt_vel", "gt_gyro"): sid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SENSOR, nm) self._gt[nm] = slice(self.model.sensor_adr[sid], self.model.sensor_adr[sid] + self.model.sensor_dim[sid]) def _settle_height(self): """Where does the robot stand in THIS task's pose? Solved kinematically, not by simulating: an unsupported robot topples during any settle long enough to converge, and you end up measuring the height of a fallen robot. """ d = mujoco.MjData(self.model) d.qpos[:3] = self.spec["root"]["pos"] for name, val in getattr(self, "init_q", {}).items(): jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) if jid >= 0: d.qpos[self.model.jnt_qposadr[jid]] = val mujoco.mj_forward(self.model, d) # lowest point of any wheel, and the radius it should be resting on lowest, radius = None, 0.0 for gi in range(self.model.ngeom): if self.model.geom_type[gi] != mujoco.mjtGeom.mjGEOM_CYLINDER: continue r = float(self.model.geom_size[gi][0]) if r < 0.03: continue z = float(d.geom_xpos[gi][2]) if lowest is None or z < lowest: lowest, radius = z, r if lowest is None: self.spawn_z = float(self.spec["root"]["pos"][2]) else: # raise or lower the body so the wheels just touch the floor self.spawn_z = float(d.qpos[2]) - (lowest - radius) self.h_target = self.spawn_z self.min_h = self.spawn_z * 0.6 def _resolve_motors(self): """Which actuators are torque-mode, and where their dof lives. This has to be re-resolved whenever the model changes. Swapping the root joint for a tether rig renumbers every dof - the free-floating robot has 14, hung from a hook it has 12 - so a cached dof index from the old model reads the wrong joint at best and runs off the end of qvel at worst. The torque-speed curve in step() dereferences these every tick, which is how a stale index crashed the whole live sim. """ import mjcf as _mjcf self.torque_idx, self.act_dof, self.no_load = [], {}, {} for i, a in enumerate(self.spec.get("actuators", [])): jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, a["joint"]) self.act_dof[i] = self.model.jnt_dofadr[jid] if a.get("mode") == "torque": lim = _mjcf.motor_limits(self.spec, a) self.torque_idx.append(i) self.no_load[i] = lim["no_load_rad_s"] if lim else 700.0 def _rebind(self): """Re-resolve everything that points into the model. Swapping the root joint (to build a tethered test rig) changes every index, so the sensor addresses and joint ids have to be looked up again or the readings silently belong to the wrong quantities. """ self.ctrl_lo = self.model.actuator_ctrlrange[:, 0].copy() self.ctrl_hi = self.model.actuator_ctrlrange[:, 1].copy() self.tau_limit = np.abs(self.model.actuator_forcerange[:, 1]).copy() self.tau_limit[self.tau_limit < 1e-9] = 1.0 self.leg_joint_ids = np.array([ self.model.jnt_qposadr[mujoco.mj_name2id( self.model, mujoco.mjtObj.mjOBJ_JOINT, n)] for n in self.init_q], dtype=int) self._resolve_motors() self.sensors = SensorBank(self.model, self.spec, self.control_dt, self.rng) self._gt = {} for nm in ("gt_quat", "gt_pos", "gt_up", "gt_vel", "gt_gyro"): sid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SENSOR, nm) self._gt[nm] = slice(self.model.sensor_adr[sid], self.model.sensor_adr[sid] + self.model.sensor_dim[sid]) self._raw = {} for nm in ("gyro", "accel", "mag"): sid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SENSOR, nm) if sid >= 0: self._raw[nm] = slice(self.model.sensor_adr[sid], self.model.sensor_adr[sid] + self.model.sensor_dim[sid]) self._base_mass = self.model.body_mass.copy() self._base_inertia = self.model.body_inertia.copy() self._base_ipos = self.model.body_ipos.copy() self._base_friction = self.model.geom_friction.copy() self._base_gain = self.model.actuator_gainprm.copy() self._base_bias = self.model.actuator_biasprm.copy() self._base_forcerange = self.model.actuator_forcerange.copy() # Tethered means the FREE JOINT WAS REPLACED, not "the model was # rebuilt". This used to be hard-coded True, so any rebuilt world # (the tilting table, scenery props) silently took the tethered # reset path: qpos zeroed, no spawn placement, a free-floating robot # left at floor height - which is the fault _place_on_table was # papering over after every table restart. self.tethered = not any( self.model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE for j in range(self.model.njnt)) # ------------------------------------------------------------------ DR def _apply_dr(self): m = self.model m.body_mass[:] = self._base_mass m.body_inertia[:] = self._base_inertia m.body_ipos[:] = self._base_ipos m.geom_friction[:] = self._base_friction m.actuator_gainprm[:] = self._base_gain m.actuator_biasprm[:] = self._base_bias m.actuator_forcerange[:] = self._base_forcerange self.latency = 0 if not self.dr.get("enabled", False): return lo, hi = self.dr.get("mass_scale", [1, 1]) scale = self.rng.uniform(lo, hi, size=m.nbody) m.body_mass[:] = self._base_mass * scale m.body_inertia[:] = self._base_inertia * scale[:, None] com = float(self.dr.get("com_offset", 0.0)) if com: root_bid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, self.spec["root"]["name"]) m.body_ipos[root_bid] = self._base_ipos[root_bid] + self.rng.uniform(-com, com, 3) lo, hi = self.dr.get("friction_scale", [1, 1]) m.geom_friction[:, 0] = self._base_friction[:, 0] * self.rng.uniform(lo, hi) lo, hi = self.dr.get("actuator_gain_scale", [1, 1]) g = self.rng.uniform(lo, hi, size=m.nu) m.actuator_gainprm[:, 0] = self._base_gain[:, 0] * g # Position servos encode kp in biasprm[1] and kv in biasprm[2]; scale # them with the gain so the servo stays consistent. m.actuator_biasprm[:, 1] = self._base_bias[:, 1] * g m.actuator_biasprm[:, 2] = self._base_bias[:, 2] * g # A stronger motor is stronger at the LIMIT too. Scaling only the # gain left MuJoCo clamping gain*ctrl to the nominal forcerange, so # the randomisation could reduce peak torque but never raise it - a # one-sided draw that trained every policy against a weaker robot # than the spec describes (expected wheel peak was ~1.30 of 1.4 Nm). m.actuator_forcerange[:] = self._base_forcerange * g[:, None] lo, hi = self.dr.get("latency_steps", [0, 0]) self.latency = int(self.rng.integers(lo, hi + 1)) # --------------------------------------------------------------- reset def reset(self): self._apply_dr() mujoco.mj_resetData(self.model, self.data) noise = dict(self.spec.get("init", {}).get("noise", {})) if getattr(self.task_obj, "init_noise", None): noise.update(self.task_obj.init_noise) root = self.spec["root"] if getattr(self, "tethered", False): self.data.qpos[:] = 0.0 # a ball joint's rest value is the identity quaternion, not zero for j in range(self.model.njnt): if self.model.jnt_type[j] == mujoco.mjtJoint.mjJNT_BALL: a = self.model.jnt_qposadr[j] self.data.qpos[a:a+4] = [1.0, 0.0, 0.0, 0.0] for name, val in self.init_q.items(): jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) self.data.qpos[self.model.jnt_qposadr[jid]] = val mujoco.mj_forward(self.model, self.data) self.sensors.reset(); self.estimator.reset() self.est_up = np.array([0.0, 0.0, 1.0]); self.est_heading = 0.0 self.step_i = 0; self.t = 0.0 self.feel = {"up_pk": 0.0, "fa_i": 0.0, "fa_pk": 0.0, "side_i": 0.0, "side_pk": 0.0, "jerk_i": 0.0, "rate_pk": 0.0} self._feel_prev_a = None self.prev_action = np.zeros(self.n_act) self.action_queue = [np.zeros(self.n_act) for _ in range(1)] self.latency = 0 self._cache_truth(); self.task_obj.reset(self, self.rng) self._next_push = 1e9 return self._obs() self.data.qpos[0:3] = root["pos"] self.data.qpos[2] = self.spawn_z + self.rng.normal(0, noise.get("base_z", 0.0)) # per-axis, so a task can randomise pitch without randomising roll base_rpy = noise.get("base_rpy", 0.0) rpy = np.array([ self.rng.normal(0, noise.get("base_roll", base_rpy)), self.rng.normal(0, noise.get("base_pitch", base_rpy)), self.rng.normal(0, noise.get("base_yaw", base_rpy)), ]) quat = np.zeros(4) mujoco.mju_euler2Quat(quat, rpy, "xyz") self.data.qpos[3:7] = quat jn = noise.get("joint", 0.0) for adr, val in zip(self.leg_joint_ids, self.q_nominal): self.data.qpos[adr] = val + self.rng.normal(0, jn) self.data.qvel[:6] = self.rng.normal(0, noise.get("base_vel", 0.0), 6) mujoco.mj_forward(self.model, self.data) self.sensors.reset() self.estimator.reset() self.est_up = np.array([0.0, 0.0, 1.0]) self.est_heading = 0.0 # a fixed but random direction for this episode's motor interference v = self.rng.normal(size=3) self._mag_axis = v / (np.linalg.norm(v) + 1e-9) self.feel = {"up_pk": 0.0, "fa_i": 0.0, "fa_pk": 0.0, "side_i": 0.0, "side_pk": 0.0, "jerk_i": 0.0, "rate_pk": 0.0} self._feel_prev_a = None self.step_i = 0 self.t = 0.0 self.prev_action = np.zeros(self.n_act) self.action_queue = [np.zeros(self.n_act) for _ in range(self.latency + 1)] self._cache_truth() self.task_obj.curriculum = self.cmd_scale self.task_obj.reset(self, self.rng) self._next_push = float(self.dr.get("push", {}).get("every_s", 1e9)) return self._obs() def _cache_truth(self): """Ground truth, refreshed once per step. Reward and scenery logic may read these; the observation may not.""" sd = self.data.sensordata self.base_pos = np.asarray(sd[self._gt["gt_pos"]], dtype=float) self.base_quat = np.asarray(sd[self._gt["gt_quat"]], dtype=float) self.vel_local = np.asarray(sd[self._gt["gt_vel"]], dtype=float) self.gyro_true = np.asarray(sd[self._gt["gt_gyro"]], dtype=float) self.up_z = float(sd[self._gt["gt_up"]][2]) self.height = float(self.base_pos[2]) def _est_update(self): """Produce the attitude the flight computer would receive.""" if not self.est_on: return np.zeros(0) if self.est_mode == "fused": q = self.base_quat yaw = math.atan2(2*(q[0]*q[3] + q[1]*q[2]), 1 - 2*(q[2]*q[2] + q[3]*q[3])) drive = float(np.mean(np.abs(self.data.actuator_force))) / 6.0 up, head = self.estimator.update( np.asarray(self.data.sensordata[self._gt["gt_up"]], dtype=float), yaw, drive) self.est_up, self.est_heading = up, head return up g = self.sensors.noisy("gyro", self.data.sensordata, self._raw) a = self.sensors.noisy("accel", self.data.sensordata, self._raw) self.estimator.update(g, a) if self.has_mag: m = self.sensors.noisy("mag", self.data.sensordata, self._raw).copy() if self.mag_coupling: # motor currents put field through the IMU; torque is the proxy drive = float(np.mean(np.abs(self.data.actuator_force))) m = m + self._mag_axis * self.mag_coupling * drive * 0.02 self.estimator.update_heading(g, m) return self.estimator.up def _feel_update(self): """THE FEELINGS: the vestibular sense, metered every tick. A kid learns to stand by feeling its balance - these are the numbers that feeling is made of, taken from the same noisy IMU the real robot will carry (never ground truth). Every episode of every stage accumulates them in env.feel; a stage prices them into its score via FEEL_WEIGHTS when its exam wants the skill smooth, not just successful. up_pk peak vertical thrust (g, above gravity) fa_i/pk fore-aft g: integral and peak (the lurch) side_i/pk lateral g: integral and peak (the wobble) jerk_i integral of accel change (the roughness) rate_pk peak body rotation rate (rad/s) """ a = self.sensors.noisy("accel", self.data.sensordata, self._raw) g = self.sensors.noisy("gyro", self.data.sensordata, self._raw) up = getattr(self, "est_up", np.array([0.0, 0.0, 1.0])) p = math.atan2(-float(up[0]), max(1e-6, float(up[2]))) # the accel publishes in sensor units (spec scale 0.1): gravity # reads ~98.1, MEASURED at standstill - never assume 9.81 _G = 98.1 ax, ay, az = float(a[0]), float(a[1]), float(a[2]) up_g = (az * math.cos(p) - ax * math.sin(p)) / _G - 1.0 fa_g = (ax * math.cos(p) + az * math.sin(p)) / _G side_g = ay / _G f = self.feel # the polish pass listens here: when FEEL_ACCUM is armed it # accumulates ROUGHNESS - tick-to-tick change, the jaggedness # of the line - NOT amplitude. A sweeping turn or a jump is a # slow wave and barely registers; chatter and micro-correction # ripple register enormously (operator's spec: smooth the # lines, don't flatten the move) global FEEL_ACCUM if FEEL_ACCUM is not None: _p = FEEL_ACCUM.get("_prev") if _p is not None: FEEL_ACCUM["rough"] = FEEL_ACCUM.get("rough", 0.0) \ + abs(fa_g - _p[0]) + abs(side_g - _p[1]) \ + 0.5 * abs(up_g - _p[2]) FEEL_ACCUM["_prev"] = (fa_g, side_g, up_g) f["up_pk"] = max(f["up_pk"], up_g) f["fa_i"] += abs(fa_g) * self.control_dt f["fa_pk"] = max(f["fa_pk"], abs(fa_g)) f["side_i"] += abs(side_g) * self.control_dt f["side_pk"] = max(f["side_pk"], abs(side_g)) prev = self._feel_prev_a if prev is not None: f["jerk_i"] += float(np.sum(np.abs(a - prev))) / 98.1 self._feel_prev_a = np.asarray(a, dtype=float).copy() f["rate_pk"] = max(f["rate_pk"], float(np.max(np.abs(g)))) def _obs(self): s = self.sensors.read(self.data.sensordata, self.t) if self.est_on: up = getattr(self, "est_up", np.array([0.0, 0.0, 1.0])) head = getattr(self, "est_heading", 0.0) est = np.concatenate([up, [math.sin(head), math.cos(head)]]) else: est = np.zeros(0) return np.concatenate([s, est, self.task_obj.obs(self), self.prev_action]).astype(np.float64) # ---------------------------------------------------------------- step def step(self, action): action = np.clip(np.asarray(action, dtype=float), -1.0, 1.0) self.action_queue.append(action) applied = self.action_queue.pop(0) full = np.zeros(len(self.act_center)) full[self.act_idx] = applied full = full * self.act_sign ctrl = np.clip(self.act_center + full * self.act_span, self.ctrl_lo, self.ctrl_hi) # A motor cannot hold peak torque at speed: back-EMF eats the voltage # headroom, so available torque falls roughly linearly to zero at the # no-load speed. Without this the wheels behave like an infinite source # and the sim flatters every fast manoeuvre. for i in self.torque_idx: w = abs(float(self.data.qvel[self.act_dof[i]])) avail = self.ctrl_hi[i] * max(0.0, 1.0 - w / self.no_load[i]) ctrl[i] = float(np.clip(ctrl[i], -avail, avail)) # write only the robot's actuators: a test rig may append its own # (the tilting platform), which the env must not zero every step self.data.ctrl[:len(ctrl)] = ctrl push = self.dr.get("push", {}) if self.dr.get("enabled") and self.t >= self._next_push: imp = self.rng.uniform(*push.get("impulse", [0, 0])) ang = self.rng.uniform(0, 2 * math.pi) self.data.qvel[0] += imp * math.cos(ang) self.data.qvel[1] += imp * math.sin(ang) self._next_push = self.t + push.get("every_s", 1e9) for _ in range(self.n_sub): mujoco.mj_step(self.model, self.data) self.t += self.control_dt self.step_i += 1 self._cache_truth() self._est_update() self._feel_update() self.task_obj.on_step(self, self.rng) reward, terms = self.task_obj.reward(self, action) fell = (self.up_z < self.min_up) or (self.height < self.min_h) \ or not np.isfinite(self.data.qpos).all() or self.task_obj.terminated(self) truncated = self.step_i >= self.max_steps self.prev_action = action info = {"terms": terms, "fell": bool(fell), "height": self.height, "up_z": self.up_z, "success": bool(self.task_obj.success(self))} return self._obs(), float(reward), bool(fell), bool(truncated), info # -------------------------------------------------------------- replay def body_names(self): return [mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, i) for i in range(1, self.model.nbody)] def frame(self): """Body world poses for the three.js bench, so the viewer needs no FK.""" f = { "t": round(self.t, 4), "pos": [[round(float(v), 5) for v in self.data.xpos[i]] for i in range(1, self.model.nbody)], "quat": [[round(float(v), 5) for v in self.data.xquat[i]] for i in range(1, self.model.nbody)], } goal = self.task_obj.goal_marker(self) if goal is not None: f["goal"] = [round(v, 4) for v in goal] return f