""" Learned-policy runtime for the RX1 (v2). The robot is driven by the LEARNED action model (models/act_model.py), not by scripted IK. A high-level planner turns the instruction into a queue of (pick cube, place target) steps; for each step the trained, goal-conditioned model outputs the arm trajectory + grasp, and that is what moves the robot: obs = [ proprio(80), pick_pos0(3), place_pos(3) ] → ActionChunkPolicy → (chunk_len, 15) [7 right-arm, 7 left-arm, grasp] → expand to the 40 MuJoCo actuators (torso/head at rest, fingers from the grasp bit) + the target cube's weld bit. Pick / place-in-box / stack / sort all use the same model — only the place target differs. If no trained checkpoint is given, this falls back to the v1 IK policy so the stack still runs. The planner here can be replaced by the Qwen2.5-VL planner (models/vlm_planner.py); it only needs to emit (cube colour, place target). """ from __future__ import annotations import logging import time from typing import Optional import numpy as np from models.base import BasePolicy from models.rx1_ik import (RX1IKSolver, load_rx1_model, discover_cube_welds, CUBE_SIZE, RIGHT_ARM_ACT, LEFT_ARM_ACT, BOX_POS, CUBE_FALLBACK) log = logging.getLogger(__name__) ARM_ACT = list(RIGHT_ARM_ACT) + list(LEFT_ARM_ACT) # 14 indices _COLORS = ("red", "blue", "green", "yellow") _STACK_KW = ("stack", "tower", "pile") _SORT_KW = ("sort", "tidy", "organi", "collect", "clean", "clear", "put away") _PICK_KW = ("pick", "grab", "grasp", "take", "get", "lift", "place", "put", "drop", "move", "fetch", "bring", "hold") _STOP_KW = ("stop", "cancel", "abort", "halt") _STEP_MIN_S = 4.0 # don't declare a step done before the motion can finish _STEP_MAX_S = 11.0 # hard timeout per step _PLACE_TOL = 0.06 # cube within this of the target ⇒ step succeeded # Open-loop playback: the model is queried ONCE per pick-place step (in-distribution # rest+goal obs) and the full predicted trajectory is played out without feedback, # which sidesteps behaviour-cloning covariate shift. _STRIDE rows are consumed per # act() call (≈ control_freq / think_freq); _RET is the lookahead window pushed. _STRIDE = 2 _RET = 16 class _LearnedRX1Policy(BasePolicy): def __init__(self, ckpt_path: str, device: str = "cpu"): from models.act_model import load_policy import torch # Match load_policy's graceful CPU degradation so inference (.to(device)) # doesn't crash when the config requests cuda on a CPU-only host. if "cuda" in str(device) and not torch.cuda.is_available(): device = "cpu" self._model, self._norm = load_policy(ckpt_path, device) self._device = device self._chunk = self._model.chunk_len model = load_rx1_model() self._ik = RX1IKSolver(model) self._nu = model.nu self._rest = self._ik.rest() self._welds = discover_cube_welds(model) self._weld_bit_by_color = {b.replace("_cube", ""): self._nu + i for i, (b, _e) in enumerate(self._welds)} self._action_dim = self._nu + len(self._welds) # task state self._task_text = "" self._queue = [] self._plan = None self._step = None # active step dict self._pick0 = None # captured cube centre when the step started self._place = None self._weld_col = None self._cur_q = None # current qpos (compliant non-arm command) self._traj = None # cached open-loop arm trajectory for the step self._cursor = 0 # playback position into self._traj self._t0 = 0.0 self._think_n = 0 log.info("[LearnedPolicy] loaded %s chunk=%d action_dim=%d device=%s", ckpt_path, self._chunk, self._action_dim, device) # planner shares the v2 grammar (stack / sort / pick-place / colours) def set_task_text(self, text: str): t = text.lower().strip() if (self._step or self._queue or self._plan) and any(k in t for k in _STOP_KW): self._queue = []; self._plan = None; self._step = None log.info("[LearnedPolicy] task cancelled") return self._task_text = text self._queue = []; self._plan = None; self._step = None colors = [c for c in _COLORS if c in t] if any(k in t for k in _STOP_KW): pass elif any(k in t for k in _STACK_KW): if len(colors) >= 2 and (" on " in t or "onto" in t or "on top" in t): self._queue = [{"color": colors[0], "kind": "on", "on": colors[1]}] else: self._plan = {"type": "stack", "colors": colors or None} elif any(k in t for k in _SORT_KW): self._plan = {"type": "sort", "colors": colors or None} elif any(k in t for k in _PICK_KW) or colors: if len(colors) >= 2 and (" on " in t or "onto" in t or "on top" in t): self._queue = [{"color": colors[0], "kind": "on", "on": colors[1]}] else: self._queue = [{"color": colors[0] if colors else None, "kind": "box"}] log.info("[LearnedPolicy] '%s' -> plan=%s queue=%d", text, self._plan, len(self._queue)) def plan_snapshot(self, objects=None): """Serializable view of the whole plan for terminal display: the parsed high-level plan, the ordered step queue (compound stack/sort plans are expanded against live object poses for a concrete preview), each step's resolved place target, and the active step's execution progress.""" # Build the display list: the active (in-flight) step first, then queued. pending = list(self._queue) if self._plan is not None and not pending and objects is not None: try: pending = self._expand(objects) except Exception: pending = [] display = ([self._step] if self._step is not None else []) + pending steps = [] for st in display: place = None try: place = [round(float(x), 3) for x in self._place_for(st, objects)] except Exception: place = None steps.append({"color": st.get("color"), "kind": st.get("kind"), "on": st.get("on"), "place": place}) def _xyz(v): return [round(float(x), 3) for x in v] if v is not None else None return { "task": self._task_text, "plan": self._plan, # {"type","colors"} or None "kind": (self._plan["type"] if self._plan else (steps[0]["kind"] if steps else None)), "steps": steps, "active": (steps[0] if self._step is not None and steps else None), "pick0": _xyz(self._pick0), "place_target": _xyz(self._place), "weld_col": self._weld_col, "cursor": int(self._cursor), "traj_len": int(len(self._traj)) if self._traj is not None else 0, "queue_remaining": len(self._queue), } def _expand(self, objects): cubes = [(k, v) for k, v in (objects or {}).items() if k != "box" and v.get("weld_bit") is not None] if not cubes: return [] def select(cols): if not cols: return list(cubes) by = {v.get("color"): (k, v) for k, v in cubes} return [by[c] for c in cols if c in by] out = [] if self._plan["type"] == "stack": sel = select(self._plan.get("colors")) if len(sel) < 2: return [] base = np.asarray(sel[0][1]["pos"], float) for lvl, (_n, info) in enumerate(sel[1:], start=1): out.append({"color": info["color"], "kind": "stack", "place": base + [0, 0, CUBE_SIZE * lvl]}) elif self._plan["type"] == "sort": for _n, info in select(self._plan.get("colors")): out.append({"color": info["color"], "kind": "box"}) return out def _place_for(self, step, objects): if step["kind"] == "stack": return np.asarray(step["place"], float) if step["kind"] == "on": base = next((v for k, v in (objects or {}).items() if k != "box" and v.get("color") == step["on"]), None) p = (np.asarray(base["pos"], float) if base is not None else CUBE_FALLBACK.get(step["on"], BOX_POS).copy()) return p + [0, 0, CUBE_SIZE] if objects and "box" in objects: return np.asarray(objects["box"]["pos"], float) + [0, 0, 0.03] return BOX_POS.copy() def _resolve_cube(self, step, objects): cubes = {k: v for k, v in (objects or {}).items() if k != "box" and v.get("weld_bit") is not None} want = step.get("color") name = None if want: name = next((k for k, v in cubes.items() if v.get("color") == want), None) if name is None and cubes: name = next(iter(cubes)) if name is None: return None v = cubes[name] return name, np.asarray(v["pos"], float), int(v["weld_bit"]) @staticmethod def _denorm_proprio_qpos(proprio): """Invert PerceptionModule._normalise_proprio for the qpos (first 40) so we can command the non-arm joints to their CURRENT pose (compliant regime — matches the expert; the arm can't reach from a firm torso).""" if proprio is None: return None p = np.asarray(proprio, np.float32).copy() if len(p) < 40: return None q = p[:40].copy() q[:3] *= 3.0 # root/torso positions were /3 q[7:40] *= 5.0 # remaining joint angles were /5 (p[3:7] unchanged) return q def _hold_chunk(self): """Idle/wait command: hold the current pose (compliant), else rest.""" row = np.zeros(self._action_dim, np.float32) row[:self._nu] = self._cur_q if self._cur_q is not None else self._rest return np.tile(row, (self._chunk, 1)) def act(self, world_embedding, task_embedding, proprio=None, text="", objects=None, grounding=None, **_): import torch self._think_n += 1 self._cur_q = self._denorm_proprio_qpos(proprio) if text and text != self._task_text: self.set_task_text(text) # start the next step if idle if self._step is None: if self._plan is not None and not self._queue: self._queue = self._expand(objects); self._plan = None if self._queue: step = self._queue[0] resolved = self._resolve_cube(step, objects) if resolved is None: return self._hold_chunk() _name, self._pick0, self._weld_col = resolved self._place = self._place_for(step, objects) self._step = step self._traj = None # force a fresh open-loop query self._cursor = 0 self._t0 = time.monotonic() log.info("[LearnedPolicy] step %s pick0=%s place=%s weld_bit=%d", step, np.round(self._pick0, 3), np.round(self._place, 3), self._weld_col) else: return self._hold_chunk() # query the LEARNED model ONCE per step → full open-loop trajectory if self._traj is None: pr = np.zeros(80, np.float32) if proprio is not None: pr[:min(80, len(proprio))] = np.asarray(proprio, np.float32)[:80] obs = np.concatenate([pr, self._pick0, self._place]).astype(np.float32) with torch.no_grad(): out = self._model(torch.from_numpy(self._norm.norm_obs(obs)) .float().unsqueeze(0).to(self._device)) self._traj = self._norm.denorm_act(out.squeeze(0).cpu().numpy()) # (T,15) self._cursor = 0 # step done when the trajectory is played out (or a hard timeout) if self._cursor >= len(self._traj) or (time.monotonic() - self._t0) > _STEP_MAX_S * 2: err = "—" if objects is not None: cur = self._resolve_cube(self._step, objects) if cur: err = round(float(np.linalg.norm(cur[1] - self._place)), 3) log.info("[LearnedPolicy] step done (played %d steps, place err=%s) -> next", self._cursor, err) self._queue.pop(0); self._step = None; self._traj = None return self._hold_chunk() # build the lookahead window from the cached trajectory (open-loop) chunk = np.zeros((_RET, self._action_dim), np.float32) for j in range(_RET): idx = min(self._cursor + j, len(self._traj) - 1) full = (self._cur_q.copy() if self._cur_q is not None else self._rest.copy()) full[ARM_ACT] = self._traj[idx, :14] grasp = self._traj[idx, 14] > 0.5 full = self._ik.set_fingers(full, closed=bool(grasp), side="right") chunk[j, :self._nu] = full if grasp and self._weld_col is not None: chunk[j, self._weld_col] = 1.0 self._cursor += _STRIDE return chunk def make_learned_policy(config: dict): """Factory: learned policy if a checkpoint is configured, else v1 IK fallback.""" ckpt = config.get("ckpt_path", "") if ckpt and __import__("os").path.exists(ckpt): return _LearnedRX1Policy(ckpt, config.get("device", "cpu")) log.warning("[LearnedPolicy] no checkpoint at %r — falling back to v1 IK policy", ckpt) from models.policy import CosmosPolicy cfg = dict(config); cfg["stub"] = True; cfg["use_rx1_ik"] = True return CosmosPolicy(cfg)