""" Shared RX1 arm IK solver and pick-and-place executor. Robot now has ALL 40 joints dynamic (nu=40): [0-2] torso (yaw, pitch, roll) [3-7] head/neck (5 DoF) [8-14] right arm (7 DoF) [15-23] right fingers (9 DoF) [24-30] left arm (7 DoF) [31-39] left fingers (9 DoF) """ from __future__ import annotations import pathlib import re import time from enum import Enum, auto from typing import List, Optional import numpy as np import mujoco ROOT = pathlib.Path(__file__).resolve().parents[1] _XML = ROOT / "rx1_mujoco" / "rx1_scene.xml" _MESHDIR = (ROOT / "rx1_mujoco" / "meshes_mjcf").as_posix() + "/" # World positions (rx1_scene.xml worldbody) — used only as a fallback when # perception is unavailable; live perceived positions win at runtime. RED_CUBE_POS = np.array([0.31, -0.05, 0.445]) BLUE_CUBE_POS = np.array([0.31, -0.10, 0.445]) GREEN_CUBE_POS = np.array([0.31, -0.15, 0.445]) YELLOW_CUBE_POS = np.array([0.31, -0.20, 0.445]) BOX_POS = np.array([0.375, -0.085, 0.50]) # Fallback spawn position per colour (keys match env _color_name output). CUBE_FALLBACK = { "red": RED_CUBE_POS, "blue": BLUE_CUBE_POS, "green": GREEN_CUBE_POS, "yellow": YELLOW_CUBE_POS, } # Cube full edge length (geom size 0.022 half) — used to compute stack heights. CUBE_SIZE = 0.044 def discover_cube_welds(model: "mujoco.MjModel") -> List[tuple]: """ Ordered ``[(body_name, eq_id), ...]`` for every weld whose welded body has a free joint — i.e. each manipulable cube. Order is the equality-constraint id (== XML declaration order), so the env and the policy derive identical weld-bit indices without sharing any state: weld bit for cube ``i`` is action index ``nu + i``. """ free_bodies = set() for j in range(model.njnt): if model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE: free_bodies.add(int(model.jnt_bodyid[j])) welds: List[tuple] = [] for e in range(model.neq): if model.eq_type[e] != mujoco.mjtEq.mjEQ_WELD: continue name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_EQUALITY, e) or "" if not name.startswith("weld_"): continue cube_bid = next((bid for bid in (int(model.eq_obj2id[e]), int(model.eq_obj1id[e])) if bid in free_bodies), None) if cube_bid is None: continue welds.append((mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, cube_bid), e)) return welds # Actuator index ranges (from build_scene.py output) RIGHT_ARM_ACT = list(range(8, 15)) # actuators 8-14 RIGHT_FINGER_ACT = list(range(15, 24)) # actuators 15-23 LEFT_ARM_ACT = list(range(24, 31)) # actuators 24-30 LEFT_FINGER_ACT = list(range(31, 40)) # actuators 31-39 TORSO_ACT = list(range(0, 3)) # actuators 0-2 HEAD_ACT = list(range(3, 8)) # actuators 3-7 def load_rx1_model() -> mujoco.MjModel: xml = _XML.read_text(encoding="utf-8") xml = re.sub(r'meshdir="[^"]*"', f'meshdir="{_MESHDIR}"', xml) return mujoco.MjModel.from_xml_string(xml) def _arm_dof_ids(model: mujoco.MjModel, joint_names: List[str]) -> List[int]: """Get velocity-DOF (Jacobian column) indices for a list of joint names.""" ids = [] for jname in joint_names: jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, jname) if jid >= 0: ids.append(int(model.jnt_dofadr[jid])) return ids def _arm_qpos_ids(model: mujoco.MjModel, joint_names: List[str]) -> List[int]: """Get qpos addresses (where each joint angle lives in qpos) for joint names.""" ids = [] for jname in joint_names: jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, jname) if jid >= 0: ids.append(int(model.jnt_qposadr[jid])) return ids _RIGHT_ARM_JOINTS = [ "right_shoul_base2shoul_joint", "right_shoul2shoul_rot_joint", "right_arm2armrot_joint", "right_armrot2elbow_joint", "right_forearm2forearmrot_joint", "right_forearmrot2forearm_pitch_joint", "right_forearm_pitch2forearm_roll_joint", ] _LEFT_ARM_JOINTS = [j.replace("right_", "left_") for j in _RIGHT_ARM_JOINTS] class RX1IKSolver: """Damped-least-squares IK for both arm grasp sites on the RX1.""" def __init__(self, model: mujoco.MjModel): self.model = model self.nu = model.nu # 40 self.nv = model.nv self._lo = model.actuator_ctrlrange[:, 0].copy() self._hi = model.actuator_ctrlrange[:, 1].copy() # Grasp site IDs self._rsite = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "right_grasp") self._lsite = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "left_grasp") # DOF columns in the Jacobian for each arm self._r_dofs = _arm_dof_ids(model, _RIGHT_ARM_JOINTS) # 7 indices self._l_dofs = _arm_dof_ids(model, _LEFT_ARM_JOINTS) # 7 indices # qpos addresses for each arm (where to write joint angles in qpos) self._r_qpos = _arm_qpos_ids(model, _RIGHT_ARM_JOINTS) self._l_qpos = _arm_qpos_ids(model, _LEFT_ARM_JOINTS) # Actuator name → index map for convenience poses self._idx = { mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i): i for i in range(model.nu) } def solve(self, target_xyz: np.ndarray, q_ctrl: np.ndarray, side: str = "right", n_iter: int = 300, lr: float = 0.5, lam: float = 0.08) -> np.ndarray: """ Damped-least-squares IK for one arm's grasp site. Returns a full ctrl vector (nu,) with only the selected arm's actuators updated; every other joint command is copied from `q_ctrl`. The solve runs in a scratch MjData by writing joint *positions* (qpos) and calling mj_kinematics — NOT by writing ctrl + mj_forward, which would never move the joints (forward dynamics integrates nothing, so the site/Jacobian would be evaluated at the rest configuration forever). """ arm_act = RIGHT_ARM_ACT if side == "right" else LEFT_ARM_ACT dof_ids = self._r_dofs if side == "right" else self._l_dofs # Jacobian cols qp_ids = self._r_qpos if side == "right" else self._l_qpos # qpos addresses site_id = self._rsite if side == "right" else self._lsite d = mujoco.MjData(self.model) q_out = q_ctrl.copy() # Seed the full robot configuration so non-arm links are posed correctly. d.qpos[:self.nu] = q_out lo = self._lo[arm_act] hi = self._hi[arm_act] q_arm = np.clip(q_out[arm_act].copy(), lo, hi) final_err = float("inf") it = 0 for it in range(n_iter): d.qpos[qp_ids] = q_arm mujoco.mj_kinematics(self.model, d) mujoco.mj_comPos(self.model, d) # fills cdof — required by mj_jacSite err = target_xyz - d.site_xpos[site_id] final_err = float(np.linalg.norm(err)) if final_err < 0.004: break jacp = np.zeros((3, self.nv)) mujoco.mj_jacSite(self.model, d, jacp, None, site_id) J = jacp[:, dof_ids] # (3, 7) — only this arm's DOFs JJT = J @ J.T + lam ** 2 * np.eye(3) dq = J.T @ np.linalg.solve(JJT, err) q_arm = np.clip(q_arm + lr * dq, lo, hi) print(f"[IK] side={side} target={np.round(target_xyz, 3)} " f"residual={final_err:.4f} m iters={it+1}/{n_iter} " f"method=jacobian_dls dof={len(dof_ids)}") q_out[arm_act] = q_arm return q_out # ── finger helpers ──────────────────────────────────────────────────────── def set_fingers(self, q: np.ndarray, closed: bool, side: str = "right") -> np.ndarray: q = q.copy() act = RIGHT_FINGER_ACT if side == "right" else LEFT_FINGER_ACT for i, idx in enumerate(act): name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, idx) val = (1.4 if "thumb" in name else 1.5) if closed else 0.0 q[idx] = val return q def set_fingers_both(self, q: np.ndarray, closed: bool) -> np.ndarray: q = self.set_fingers(q, closed, "right") q = self.set_fingers(q, closed, "left") return q # ── convenience poses (full ctrl vectors) ───────────────────────────────── def _set(self, q: np.ndarray, **kw: float) -> np.ndarray: q = q.copy() for name, val in kw.items(): if name in self._idx: q[self._idx[name]] = val return q def rest(self) -> np.ndarray: q = np.zeros(self.nu) q = self._set(q, act_right_shoul2shoul_rot_joint=0.4, act_right_armrot2elbow_joint=-0.1, act_left_shoul2shoul_rot_joint=0.4, act_left_armrot2elbow_joint=-0.1, # tilt the head down so the ego camera sees the tabletop # workspace — required for camera-based perception. act_neck_yaw2pitch_joint=0.5, act_neck_pitch2head_depth_cam_mount_joint=0.3) return self.set_fingers_both(q, closed=False) def raise_arms(self) -> np.ndarray: q = np.zeros(self.nu) q = self._set(q, act_right_shoul2shoul_rot_joint=-1.57, act_left_shoul2shoul_rot_joint=-1.57) return self.set_fingers_both(q, closed=False) def wave(self) -> np.ndarray: q = np.zeros(self.nu) q = self._set(q, act_right_shoul_base2shoul_joint=-0.4, act_right_shoul2shoul_rot_joint=-1.1, act_right_armrot2elbow_joint=-0.9, act_right_forearm2forearmrot_joint=0.5, act_left_shoul2shoul_rot_joint=0.4, act_left_armrot2elbow_joint=-0.1) return self.set_fingers_both(q, closed=False) def look_down(self, q: np.ndarray) -> np.ndarray: """Tilt head/neck toward the table.""" return self._set(q, act_neck_yaw2pitch_joint=0.4, act_neck_pitch2head_depth_cam_mount_joint=0.3) # ── pick-and-place state machine ────────────────────────────────────────────── class Phase(Enum): IDLE = auto() PREGRASP = auto() LOWER = auto() CLOSE = auto() WELD = auto() LIFT = auto() CARRY = auto() PLACE = auto() RELEASE = auto() RETRACT = auto() _DUR = { Phase.PREGRASP: 3.5, Phase.LOWER: 2.5, Phase.CLOSE: 1.2, Phase.WELD: 0.3, Phase.LIFT: 3.0, Phase.CARRY: 3.0, Phase.PLACE: 2.5, Phase.RELEASE: 1.0, Phase.RETRACT: 2.5, } class PickExecutor: """ State-machine pick-and-place using right-arm Jacobian IK + weld constraint. All 40 actuators are tracked; non-arm joints are held at their rest pose. """ def __init__(self, ik: RX1IKSolver): self._ik = ik self.phase = Phase.IDLE self._t = 0.0 self._ctrl = ik.rest() self.weld_active = False self.weld_eq_id = -1 self._cube = RED_CUBE_POS.copy() self._box = BOX_POS.copy() self._approach = None # (n,3) planner waypoints, current→pregrasp @property def running(self) -> bool: return self.phase != Phase.IDLE @property def target(self) -> np.ndarray: return self._ctrl.copy() def start(self, cube_pos: np.ndarray, eq_id: int, box_pos: Optional[np.ndarray] = None, approach_path: Optional[np.ndarray] = None): """ Begin a pick. `cube_pos`/`box_pos` are LIVE perceived positions (not constants); `approach_path` is an optional (n,3) set of planner waypoints the PREGRASP phase follows for a smooth, table-clearing arc. """ if self.phase != Phase.IDLE: return self._cube = np.asarray(cube_pos, dtype=float).copy() self._box = (np.asarray(box_pos, dtype=float).copy() if box_pos is not None else BOX_POS.copy()) self._approach = (np.asarray(approach_path, dtype=float) if approach_path is not None else None) self.weld_eq_id = eq_id self.weld_active = False self._ctrl = self._ik.rest() self.phase = Phase.PREGRASP self._t = time.monotonic() plan_n = 0 if self._approach is None else len(self._approach) print(f"[PickExecutor] START cube={np.round(self._cube, 3)} " f"box={np.round(self._box, 3)} eq_id={eq_id} " f"approach_waypoints={plan_n}") def reset(self): self.phase = Phase.IDLE self.weld_active = False def step(self, q_cur: np.ndarray) -> np.ndarray: if self.phase == Phase.IDLE: return self._ctrl.copy() elapsed = time.monotonic() - self._t if self.phase == Phase.PREGRASP: pre = self._cube + np.array([0.0, 0.0, 0.12]) # Follow the diffusion planner's approach arc if one was supplied: # interpolate along the waypoints by phase progress so the hand # sweeps a smooth, table-clearing path to the pregrasp point. if self._approach is not None and len(self._approach) > 1: frac = min(1.0, elapsed / _DUR[Phase.PREGRASP]) idx = frac * (len(self._approach) - 1) lo = int(np.floor(idx)); hi = min(lo + 1, len(self._approach) - 1) wp = self._approach[lo] + (idx - lo) * (self._approach[hi] - self._approach[lo]) pre = wp self._ctrl = self._ik.set_fingers( self._ik.solve(pre, q_cur, side="right"), closed=False) if elapsed > _DUR[Phase.PREGRASP]: self._go(Phase.LOWER) elif self.phase == Phase.LOWER: pos = self._cube + np.array([0.0, 0.0, 0.02]) self._ctrl = self._ik.set_fingers( self._ik.solve(pos, q_cur, side="right"), closed=False) if elapsed > _DUR[Phase.LOWER]: self._go(Phase.CLOSE) elif self.phase == Phase.CLOSE: self._ctrl = self._ik.set_fingers(q_cur.copy(), closed=True) if elapsed > _DUR[Phase.CLOSE]: self._go(Phase.WELD) elif self.phase == Phase.WELD: self.weld_active = True if elapsed > _DUR[Phase.WELD]: self._go(Phase.LIFT) elif self.phase == Phase.LIFT: lift = self._cube + np.array([0.0, 0.0, 0.32]) self._ctrl = self._ik.set_fingers( self._ik.solve(lift, q_cur, side="right"), closed=True) if elapsed > _DUR[Phase.LIFT]: self._go(Phase.CARRY) elif self.phase == Phase.CARRY: if elapsed > _DUR[Phase.CARRY]: self._go(Phase.PLACE) elif self.phase == Phase.PLACE: self._ctrl = self._ik.set_fingers( self._ik.solve(self._box, q_cur, side="right"), closed=True) if elapsed > _DUR[Phase.PLACE]: self._go(Phase.RELEASE) elif self.phase == Phase.RELEASE: self._ctrl = self._ik.set_fingers(q_cur.copy(), closed=False) self.weld_active = False if elapsed > _DUR[Phase.RELEASE]: self._go(Phase.RETRACT) elif self.phase == Phase.RETRACT: self._ctrl = self._ik.rest() if elapsed > _DUR[Phase.RETRACT]: self.phase = Phase.IDLE print("[PickExecutor] pick-and-place complete.") return self._ctrl.copy() def _go(self, ph: Phase): elapsed = time.monotonic() - self._t print(f"[PickExecutor] phase {self.phase.name} -> {ph.name} " f"(held {elapsed:.2f}s) cube={np.round(self._cube, 3)}") self.phase = ph self._t = time.monotonic()