#!/usr/bin/env python3 """ RX1 Robot MuJoCo Viewer — IK-based cube pick-and-place + brain integration. Uses damped-least-squares Jacobian IK so the hand actually reaches the cube, then activates the weld constraint (rx1_scene.xml has weld_red_cube / weld_blue_cube) to attach and lift it. Non-pick tasks use smooth joint-space interpolation. Usage ----- python scripts/rx1_viewer.py [--port 8788] python scripts/rx1_viewer.py --no-brain --demo # standalone cycle """ from __future__ import annotations import argparse import pathlib import re import sys import threading import time import urllib.request import json as _json from enum import Enum, auto import numpy as np import mujoco import mujoco.viewer # ── paths ──────────────────────────────────────────────────────────────────── ROOT = pathlib.Path(__file__).resolve().parents[1] XML_PATH = ROOT / "rx1_mujoco" / "rx1_scene.xml" MESHDIR = (ROOT / "rx1_mujoco" / "meshes_mjcf").as_posix() + "/" _xml = XML_PATH.read_text(encoding="utf-8") _xml = re.sub(r'meshdir="[^"]*"', f'meshdir="{MESHDIR}"', _xml) model = mujoco.MjModel.from_xml_string(_xml) data = mujoco.MjData(model) _nu = model.nu # 17 position actuators _nv = model.nv # 29 velocity DOF (17 arm + 6 red_cube + 6 blue_cube) # ── MuJoCo IDs ─────────────────────────────────────────────────────────────── _GRASP_SITE = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "right_grasp") _EQ_RED = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, "weld_red_cube") _EQ_BLUE = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, "weld_blue_cube") # cube world positions (from rx1_scene.xml worldbody) _RED_POS = np.array([0.31, -0.05, 0.445]) _BLUE_POS = np.array([0.31, -0.10, 0.445]) _BOX_POS = np.array([0.375, -0.085, 0.50]) # above the open box # ── actuator helpers ────────────────────────────────────────────────────────── _ACT = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(_nu)] _IDX = {n: i for i, n in enumerate(_ACT)} _LO = model.actuator_ctrlrange[:, 0].copy() _HI = model.actuator_ctrlrange[:, 1].copy() def _p(**kw) -> np.ndarray: v = np.zeros(_nu) for k, val in kw.items(): if k in _IDX: v[_IDX[k]] = val return v # ── finger helpers ──────────────────────────────────────────────────────────── _FIN = ["act_right_finger0","act_right_finger0_tip", "act_right_finger1","act_right_finger1_tip", "act_right_finger2","act_right_finger2_tip", "act_right_finger3","act_right_finger3_tip", "act_right_thumb"] def _set_fingers(q: np.ndarray, closed: bool) -> np.ndarray: q = q.copy() for f in _FIN: if f not in _IDX: continue if "thumb" in f: q[_IDX[f]] = 1.4 if closed else 0.0 else: q[_IDX[f]] = 1.5 if closed else 0.0 return q # ── Jacobian IK ─────────────────────────────────────────────────────────────── def _ik(target: np.ndarray, q_init: np.ndarray, n_iter: int = 250, lr: float = 0.25, lam: float = 0.04) -> np.ndarray: """ Damped-least-squares IK for the right_grasp site. Works in a scratch MjData so the live simulation is untouched. Only the first _nu (arm) DOFs are optimised; cube DOFs stay at rest. """ d = mujoco.MjData(model) q = np.clip(q_init.copy(), _LO, _HI) for _ in range(n_iter): d.qpos[:_nu] = q mujoco.mj_kinematics(model, d) err = target - d.site_xpos[_GRASP_SITE] if np.linalg.norm(err) < 0.004: break jacp = np.zeros((3, _nv)) mujoco.mj_jacSite(model, d, jacp, None, _GRASP_SITE) J = jacp[:, :_nu] # arm DOFs only JJT = J @ J.T + lam ** 2 * np.eye(3) dq = J.T @ np.linalg.solve(JJT, err) q = np.clip(q + lr * dq, _LO, _HI) return q # ── static arm poses ────────────────────────────────────────────────────────── P_REST = _set_fingers(_p(act_right_shoul2shoul_rot_joint=0.4, act_right_armrot2elbow_joint=-0.1), closed=False) P_RAISE = _set_fingers(_p(act_right_shoul2shoul_rot_joint=-1.57, act_right_armrot2elbow_joint=0.0), closed=False) P_WAVE = _set_fingers(_p(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), closed=False) P_REACH = _set_fingers(_p(act_right_shoul_base2shoul_joint=0.2, act_right_shoul2shoul_rot_joint=-0.85, act_right_arm2armrot_joint=0.3, act_right_armrot2elbow_joint=-0.6, act_right_forearmrot2forearm_pitch_joint=-0.2), closed=False) P_WALK = _set_fingers(_p(act_torso_yaw2pitch_joint=0.18, act_right_shoul2shoul_rot_joint=-0.3, act_right_armrot2elbow_joint=-0.2), closed=False) # ── pick-and-place executor ─────────────────────────────────────────────────── class Phase(Enum): IDLE = auto() PREGRASP = auto() # IK above cube LOWER = auto() # IK to cube CLOSE = auto() # close fingers WELD = auto() # activate weld constraint LIFT = auto() # IK upward CARRY = auto() # hold PLACE = auto() # IK to box RELEASE = auto() # open fingers + deactivate weld RETRACT = auto() # return to rest _PHASE_DUR = { Phase.PREGRASP: 3.0, 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, } _PHASE_SEQ = [Phase.PREGRASP, Phase.LOWER, Phase.CLOSE, Phase.WELD, Phase.LIFT, Phase.CARRY, Phase.PLACE, Phase.RELEASE, Phase.RETRACT] class PickExecutor: def __init__(self): self.phase = Phase.IDLE self._t = 0.0 self._ctrl_tgt = P_REST.copy() self._weld_on = False self._weld_id = -1 self._cube_pos = _RED_POS.copy() # ── public ─────────────────────────────────────────────────────────────── def start(self, cube: str = "red"): if self.phase != Phase.IDLE: return self._cube_pos = _RED_POS if cube == "red" else _BLUE_POS self._weld_id = _EQ_RED if cube == "red" else _EQ_BLUE self.phase = Phase.PREGRASP self._t = time.monotonic() print(f"[rx1_viewer] pick → {cube} cube at {self._cube_pos}") @property def target(self) -> np.ndarray: return self._ctrl_tgt.copy() @property def weld_active(self) -> bool: return self._weld_on @property def weld_id(self) -> int: return self._weld_id @property def running(self) -> bool: return self.phase != Phase.IDLE def step(self, q_cur: np.ndarray): """Call once per viewer frame to advance state machine.""" if self.phase == Phase.IDLE: return elapsed = time.monotonic() - self._t if self.phase == Phase.PREGRASP: pre = self._cube_pos + np.array([0.0, 0.0, 0.12]) self._ctrl_tgt = _set_fingers(_ik(pre, q_cur), closed=False) if elapsed > _PHASE_DUR[Phase.PREGRASP]: self._next(Phase.LOWER) elif self.phase == Phase.LOWER: pos = self._cube_pos + np.array([0.0, 0.0, 0.018]) self._ctrl_tgt = _set_fingers(_ik(pos, q_cur), closed=False) if elapsed > _PHASE_DUR[Phase.LOWER]: self._next(Phase.CLOSE) elif self.phase == Phase.CLOSE: self._ctrl_tgt = _set_fingers(q_cur, closed=True) if elapsed > _PHASE_DUR[Phase.CLOSE]: self._next(Phase.WELD) elif self.phase == Phase.WELD: self._weld_on = True if elapsed > _PHASE_DUR[Phase.WELD]: self._next(Phase.LIFT) elif self.phase == Phase.LIFT: lift = self._cube_pos + np.array([0.0, 0.0, 0.32]) self._ctrl_tgt = _set_fingers(_ik(lift, q_cur), closed=True) if elapsed > _PHASE_DUR[Phase.LIFT]: self._next(Phase.CARRY) elif self.phase == Phase.CARRY: # hold — just keep current target if elapsed > _PHASE_DUR[Phase.CARRY]: self._next(Phase.PLACE) elif self.phase == Phase.PLACE: self._ctrl_tgt = _set_fingers(_ik(_BOX_POS, q_cur), closed=True) if elapsed > _PHASE_DUR[Phase.PLACE]: self._next(Phase.RELEASE) elif self.phase == Phase.RELEASE: self._ctrl_tgt = _set_fingers(q_cur, closed=False) self._weld_on = False if elapsed > _PHASE_DUR[Phase.RELEASE]: self._next(Phase.RETRACT) elif self.phase == Phase.RETRACT: self._ctrl_tgt = P_REST.copy() if elapsed > _PHASE_DUR[Phase.RETRACT]: self.phase = Phase.IDLE print("[rx1_viewer] pick-and-place complete.") def _next(self, ph: Phase): print(f"[rx1_viewer] pick phase → {ph.name}") self.phase = ph self._t = time.monotonic() # ── brain polling ───────────────────────────────────────────────────────────── _target = P_REST.copy() _task_str = "" _is_wave = False _is_pick = False _pick_cube= "red" _lock = threading.Lock() def _get_status(port: int) -> dict: try: url = f"http://localhost:{port}/status" with urllib.request.urlopen(url, timeout=1) as r: return _json.loads(r.read()) except Exception: return {} def _task_to_pose(task: str): """Returns (pose, is_wave, is_pick, cube_color).""" t = task.lower() if any(w in t for w in ["raise","above head","lift arm","arms up","overhead"]): return P_RAISE, False, False, "red" if any(w in t for w in ["wave","greet","hello"]): return P_WAVE, True, False, "red" if any(w in t for w in ["reach","extend"]): return P_REACH, False, False, "red" if any(w in t for w in ["walk","forward","move","step"]): return P_WALK, False, False, "red" cube = "blue" if "blue" in t else "red" if any(w in t for w in ["grab","grasp","grip","pick","take","retrieve","cube","object"]): return P_REACH, False, True, cube # pose hint unused — executor runs IK return P_REST, False, False, "red" def _poll_brain(port: int): global _target, _task_str, _is_wave, _is_pick, _pick_cube while True: s = _get_status(port) task = s.get("task", "") if task and task != _task_str: pose, wave, pick, cube = _task_to_pose(task) with _lock: _target = pose _task_str = task _is_wave = wave _is_pick = pick _pick_cube = cube print(f"[rx1_viewer] ← brain: '{task}' pick={pick} wave={wave}") time.sleep(0.5) # ── main ────────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=8788) parser.add_argument("--task", default="") parser.add_argument("--no-brain", action="store_true") parser.add_argument("--demo", action="store_true", help="Auto-cycle poses") args = parser.parse_args() global _target, _task_str, _is_wave, _is_pick, _pick_cube print(f"[rx1_viewer] model nu={_nu} nv={_nv} grasp_site={_GRASP_SITE}") if args.task: pose, wave, pick, cube = _task_to_pose(args.task) _target, _is_wave, _is_pick, _pick_cube = pose, wave, pick, cube _task_str = args.task if not args.no_brain: threading.Thread(target=_poll_brain, args=(args.port,), daemon=True).start() print(f"[rx1_viewer] polling brain on port {args.port}") # init physics mujoco.mj_resetData(model, data) mujoco.mj_forward(model, data) data.eq_active[:] = 0 # all welds off at start current = np.zeros(_nu) alpha = 0.04 exec_ = PickExecutor() _DEMO = [("rest",P_REST,False,False,"red"), ("raise",P_RAISE,False,False,"red"), ("pick red cube",P_REACH,False,True,"red"), ("wave",P_WAVE,True,False,"red"), ("pick blue cube",P_REACH,False,True,"blue")] _di, _dt = 0, time.monotonic() with mujoco.viewer.launch_passive(model, data) as v: v.cam.distance = 2.0 v.cam.elevation = -20 v.cam.azimuth = 150 v.cam.lookat[:] = [0.12, -0.05, 0.85] print("[rx1_viewer] viewer open") while v.is_running(): with v.lock(): # demo mode if args.demo and args.no_brain and not exec_.running: if time.monotonic() - _dt > 9.0: _di = (_di + 1) % len(_DEMO) name, pose, wave, pick, cube = _DEMO[_di] with _lock: _target, _is_wave, _is_pick, _pick_cube = pose, wave, pick, cube _dt = time.monotonic() print(f"[rx1_viewer] demo → '{name}'") # latch pick intent → executor with _lock: tgt = _target.copy() wave = _is_wave pick = _is_pick cube = _pick_cube if pick and not exec_.running: exec_.start(cube) with _lock: _is_pick = False # consume the intent # executor drives target during pick if exec_.running: exec_.step(current) tgt = exec_.target wave = False # weld constraint sync if exec_.running: if _EQ_RED >= 0: data.eq_active[_EQ_RED] = 1 if (exec_.weld_active and exec_.weld_id == _EQ_RED) else 0 if _EQ_BLUE >= 0: data.eq_active[_EQ_BLUE] = 1 if (exec_.weld_active and exec_.weld_id == _EQ_BLUE) else 0 else: if _EQ_RED >= 0: data.eq_active[_EQ_RED] = 0 if _EQ_BLUE >= 0: data.eq_active[_EQ_BLUE] = 0 # wave animation if wave: t_ = time.monotonic() tgt[_IDX.get("act_right_forearm2forearmrot_joint", 5)] += 0.55 * np.sin(3.0 * t_) tgt[_IDX.get("act_right_forearmrot2forearm_pitch_joint", 6)] += 0.28 * np.sin(3.0 * t_ + 1.0) # smooth interpolation current[:] = current + alpha * (tgt - current) data.ctrl[:] = current mujoco.mj_step(model, data) v.sync() time.sleep(0.002) print("[rx1_viewer] closed.") if __name__ == "__main__": main()