""" Cosmos Policy wrapper — Cosmos-Policy-*-Predict2-2B. Architecture ------------ Cosmos Policy is Cosmos-Predict2-2B-Video2World fine-tuned for robot control. Actions, proprioception, and value estimates are encoded as latent frames and jointly denoised with future video in the same diffusion process. Inference API (cosmos-policy package) -------------------------------------- from cosmos_policy.experiments.robot.cosmos_utils import get_model, get_action cfg = PolicyEvalConfig( config="cosmos_predict2_2b_480p_libero__inference_only", ckpt_path="nvidia/Cosmos-Policy-LIBERO-Predict2-2B", ) model, cosmos_config = get_model(cfg) action_dict = get_action( cfg, model, dataset_stats, observation, task_description ) actions = action_dict["actions"] # (chunk_len, action_dim) Observation format expected by cosmos-policy --------------------------------------------- { "rgb_static": (H, W, 3) uint8 — overview/third-person camera "rgb_gripper": (H, W, 3) uint8 — right wrist camera "robot_obs": (proprio_dim,) float32 } (exact keys vary per checkpoint — see each checkpoint's README) Installation ------------ git clone https://github.com/NVlabs/cosmos-policy pip install -e cosmos-policy huggingface-cli login # accept NVIDIA Noncommercial License Action space gap ---------------- Cosmos-Policy-LIBERO outputs 7-DoF arm actions. This 21-DoF humanoid requires fine-tuning the action head: 1. Collect humanoid demonstrations in MuJoCo. 2. Fine-tune: https://github.com/NVlabs/cosmos-policy (training scripts). 3. Set action_dim: 21 in configs/models.yaml. Until then, the first 7 actions are copied and the rest are zero-padded. """ from __future__ import annotations import logging import re import time from typing import Optional import numpy as np from models.base import BasePolicy log = logging.getLogger(__name__) # Config name used by cosmos-policy for the LIBERO checkpoint _LIBERO_CONFIG = "cosmos_predict2_2b_480p_libero__inference_only" _ROBOCASA_CONFIG = "cosmos_predict2_2b_480p_robocasa__inference_only" # ─────────────────────────────────────────────────────────────── stub class _StubPolicy(BasePolicy): def __init__(self, chunk_len: int, action_dim: int, noise: float, rng: np.random.Generator): self._chunk_len = chunk_len self._action_dim = action_dim self._noise = noise self._rng = rng def act(self, world_embedding: np.ndarray, task_embedding: np.ndarray, **_) -> np.ndarray: return self._rng.normal(0, self._noise, (self._chunk_len, self._action_dim)).astype(np.float32) # ─────────────────────────────────────────────────────────────── real class _RealPolicy(BasePolicy): """Wraps cosmos-policy's get_model / get_action API.""" def __init__(self, model_id: str, device: str, dtype_str: str, cfg: dict): import torch import json import pickle from huggingface_hub import hf_hub_download try: from cosmos_policy.experiments.robot.cosmos_utils import get_model from cosmos_policy.experiments.robot.libero.run_libero_eval import PolicyEvalConfig except ImportError as e: raise RuntimeError( "cosmos-policy not found. Install with:\n" " git clone https://github.com/NVlabs/cosmos-policy\n" " pip install -e cosmos-policy\n" f"Original error: {e}" ) from e self._device = device self._dtype = getattr(torch, dtype_str, torch.bfloat16) self._chunk_len = cfg.get("action_chunk_len", 10) self._action_dim = cfg.get("action_dim", 21) # Resolve cosmos-policy config name from the model_id if "libero" in model_id.lower(): config_name = _LIBERO_CONFIG elif "robocasa" in model_id.lower(): config_name = _ROBOCASA_CONFIG else: config_name = _LIBERO_CONFIG log.warning("Unknown policy checkpoint; defaulting to LIBERO config.") self._eval_cfg = PolicyEvalConfig( config=config_name, ckpt_path=model_id, ) log.info("Loading policy %s ...", model_id) self._model, self._cosmos_cfg = get_model(self._eval_cfg) log.info("Policy loaded.") # Dataset statistics are bundled with each checkpoint on HF Hub stats_file = hf_hub_download(repo_id=model_id, filename="libero_dataset_statistics.json") with open(stats_file) as f: self._dataset_stats = json.load(f) def act( self, world_embedding: np.ndarray, task_embedding: np.ndarray, images: Optional[dict] = None, proprio: Optional[np.ndarray] = None, text: str = "", **_, ) -> np.ndarray: """ Parameters ---------- world_embedding : (latent_dim,) — informational; cosmos-policy re-encodes context internally from the images task_embedding : (embedding_dim,) — not used directly; text is used images : {"overview": uint8, "right_wrist": uint8, "left_wrist": uint8} proprio : (proprio_dim,) float32 text : natural-language task string """ from cosmos_policy.experiments.robot.cosmos_utils import get_action # Build observation dict in cosmos-policy format observation = {} if images: observation["rgb_static"] = images.get("overview", _blank()) observation["rgb_gripper"] = images.get("right_wrist", _blank()) if proprio is not None: observation["robot_obs"] = proprio.astype(np.float32) action_dict = get_action( self._eval_cfg, self._model, self._dataset_stats, observation, text, ) raw = np.array(action_dict["actions"], dtype=np.float32) # (chunk, native_dim) # Adapt to humanoid action dimension native_dim = raw.shape[-1] chunk = min(raw.shape[0], self._chunk_len) out = np.zeros((self._chunk_len, self._action_dim), dtype=np.float32) copy_dim = min(native_dim, self._action_dim) out[:chunk, :copy_dim] = raw[:chunk, :copy_dim] return out def _blank(h: int = 224, w: int = 224) -> np.ndarray: return np.zeros((h, w, 3), dtype=np.uint8) # ─────────────────────────────────────────────────────────────── rx1 ik policy class _RX1IKPolicy(BasePolicy): """ VLA policy for the fully-dynamic RX1 (nu=40 joints + one weld bit per cube). Pipeline per think cycle: 1. PERCEIVE — parse task text into a multi-step plan; measure distances 2. PLAN — pop the next pick→place step, bind target to a perceived cube 3. ACT — PickExecutor state machine drives Jacobian IK targets 4. LOG — print per-cycle telemetry: phase, site-to-target dist, method Task vocabulary (free-form natural language, any cube colour): • pick / place / put / move [in box | on ] • stack [the cubes | …] — build a tower • sort / tidy / clear / collect — place every cube into the box • wave · raise arms · look · rest — gestures / postures • stop / cancel — abort the running task Higher-level tasks (stack, sort, multi-cube) are decomposed into a queue of single pick→place steps; each step reuses the proven PickExecutor primitive, so adding behaviours never touches the low-level motion code. Action layout (chunk_len, nu + n_cubes): [:, 0:nu] all 40 joint position targets (rad) [:, nu + i] weld bit for cube i (discover_cube_welds order) """ _COLORS = ("red", "blue", "green", "yellow") _PICK_KW = ("pick", "grab", "grasp", "take", "get", "lift", "retrieve", "fetch", "bring", "place", "put", "drop", "move", "hold") _STACK_KW = ("stack", "tower", "pile") _SORT_KW = ("sort", "tidy", "organi", "collect", "clean", "clear", "put away") _RAISE_KW = ("raise", "arms up", "high", "hands up") _WAVE_KW = ("wave", "hello", "hi", "greet") _LOOK_KW = ("look", "head", "face") _STOP_KW = ("stop", "cancel", "abort", "halt") # Empirical grasp-site → cube-centre vertical offset (the IK targets the hand # site; the welded cube hangs ~this far below it). Used so place/stack targets # land the cube, not the palm, at the intended height. Matches the legacy box # drop (BOX_POS sits ~this far above the box floor). _CARRY_Z = 0.055 def __init__(self, chunk_len: int, action_dim: int): from models.rx1_ik import (load_rx1_model, RX1IKSolver, PickExecutor, discover_cube_welds) model = load_rx1_model() self._ik = RX1IKSolver(model) self._executor = PickExecutor(self._ik) self._chunk_len = chunk_len self._nu = model.nu # 40 self._welds = discover_cube_welds(model) # [(body, eq_id)] self._action_dim = self._nu + len(self._welds) # 40 + n_cubes (override cfg) # colour → weld-bit index (nu + position), same ordering the env uses. self._weld_bit_by_color = {} for i, (body, _eq) in enumerate(self._welds): col = body.replace("_cube", "") self._weld_bit_by_color[col] = self._nu + i self._task_text = "" self._intent = "rest" self._queue = [] # list of pick→place step dicts self._plan = None # unexpanded high-level task {"type",...} self._weld_col = None # weld bit of the cube currently grasped self._cur_step_id = None # id() of the front step (for wait timer) self._step_since = 0.0 # when the front step became active self._think_n = 0 self._planner = None # diffusion path planner (injected by Brain) log.info("[RX1IKPolicy] init: nu=%d cubes=%d action_dim=%d method=jacobian_dls", self._nu, len(self._welds), self._action_dim) def set_path_planner(self, planner): """Brain injects the diffusion planner so the pick can follow a planned approach arc rather than a straight line to the cube.""" self._planner = planner # ── helpers: colours, target binding ────────────────────────────────────── def _colors_in(self, t: str) -> list: """Colours mentioned in `t`, in order of first appearance (deduped).""" seen, out = set(), [] for w in t.replace(",", " ").split(): if w in self._COLORS and w not in seen: seen.add(w); out.append(w) return out def _resolve_target(self, objects, want): """ Bind a (possibly None) colour to a PERCEIVED cube → live 3-D goal. Returns (name, pos, weld_bit) or None. Colour match wins; else name substring; else the first available cube. Position is the perceived (camera or sim-state) value — never a hardcoded constant. """ if not objects: return None cubes = {k: v for k, v in objects.items() if k != "box" and v.get("weld_bit") is not None} if not cubes: return None match = None if want: match = next((k for k, v in cubes.items() if v.get("color") == want), None) if match is None: match = next((k for k in cubes if want in k), None) if match is None: match = next(iter(cubes)) v = cubes[match] return match, np.asarray(v["pos"], dtype=float), int(v["weld_bit"]) def _place_pos_for_step(self, step, objects): """Resolve where this step releases the cube (grasp-site target).""" from models.rx1_ik import BOX_POS, CUBE_FALLBACK, CUBE_SIZE kind = step.get("place_kind", "box") if kind == "stack": return np.asarray(step["place_pos"], dtype=float) if kind == "on": onc = step.get("on_color") base = None if objects: base = next((v for k, v in objects.items() if k != "box" and v.get("color") == onc), None) base_pos = (np.asarray(base["pos"], dtype=float) if base is not None else CUBE_FALLBACK.get(onc, BOX_POS).copy()) p = base_pos.copy() p[2] += CUBE_SIZE + self._CARRY_Z # release one cube-height above return p # default: into the open box if objects and "box" in objects: return np.asarray(objects["box"]["pos"], dtype=float) + np.array([0.0, 0.0, 0.07]) return BOX_POS.copy() # ── intent parsing (PERCEIVE layer) ────────────────────────────────────── def set_task_text(self, text: str): t = text.lower().strip() # Cancel a running task on request; otherwise let it finish. if self._executor.running or self._queue or self._plan is not None: if any(k in t for k in self._STOP_KW): log.info("[PERCEIVE] '%s' -> cancel stopping task", text) self._executor.reset(); self._queue = []; self._plan = None self._intent = "rest" return self._task_text = text self._queue = [] self._plan = None self._intent = "rest" colors = self._colors_in(t) if any(k in t for k in self._STOP_KW): self._intent = "rest" elif any(k in t for k in self._STACK_KW): if len(colors) >= 2 and (" on " in t or "onto" in t or "on top" in t): self._queue = [{"color": colors[0], "place_kind": "on", "on_color": colors[1], "label": f"stack {colors[0]} on {colors[1]}"}] else: self._plan = {"type": "stack", "colors": colors or None} elif any(k in t for k in self._SORT_KW): self._plan = {"type": "sort", "colors": colors or None} elif any(k in t for k in self._PICK_KW) or colors: color = colors[0] if colors else None if len(colors) >= 2 and (" on " in t or "onto" in t or "on top" in t): self._queue = [{"color": color, "place_kind": "on", "on_color": colors[1], "label": f"place {color} on {colors[1]}"}] else: self._queue = [{"color": color, "place_kind": "box", "label": f"pick {color or 'cube'} → box"}] elif any(k in t for k in self._WAVE_KW): self._intent = "wave" elif any(k in t for k in self._RAISE_KW): self._intent = "raise_arms" elif any(k in t for k in self._LOOK_KW): self._intent = "look_down" else: self._intent = "rest" if self._plan or self._queue: self._step_since = time.monotonic() self._cur_step_id = None kind = ("plan:" + self._plan["type"]) if self._plan else ( f"queue[{len(self._queue)}]" if self._queue else self._intent) log.info("[PERCEIVE] text='%s' -> %s (colours=%s)", text, kind, colors) def _expand_plan(self, objects) -> list: """Turn a high-level plan (stack/sort) into a queue of pick→place steps, using live perception to know which cubes exist and where they are.""" from models.rx1_ik import CUBE_SIZE 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(colors): if not colors: return list(cubes) by_col = {v.get("color"): (k, v) for k, v in cubes} return [by_col[c] for c in colors if c in by_col] ptype = self._plan["type"] sel = select(self._plan.get("colors")) steps = [] if ptype == "stack": if len(sel) < 2: log.info("[PLAN] stack needs >=2 cubes, have %d — nothing to do", len(sel)) return [] base_name, base_info = sel[0] base = np.asarray(base_info["pos"], dtype=float) for level, (name, info) in enumerate(sel[1:], start=1): place = base.copy() place[2] = base[2] + CUBE_SIZE * level + self._CARRY_Z steps.append({"color": info["color"], "place_kind": "stack", "place_pos": place, "label": f"stack {info['color']} on {base_info['color']} (L{level})"}) log.info("[PLAN] stack → base=%s, %d cube(s) on top", base_info["color"], len(steps)) elif ptype == "sort": for name, info in sel: steps.append({"color": info["color"], "place_kind": "box", "label": f"sort {info['color']} → box"}) log.info("[PLAN] sort → %d cube(s) into the box", len(steps)) return steps def _start_step(self, step, q_cur, hand_now, objects, grounding) -> str: """Try to begin `step`. Returns 'wait' (vision not locked yet), 'started' (executor launched) or 'skip' (couldn't resolve a target).""" want = step.get("color") # look-then-grasp: wait until the ego camera actually sees the target # colour (any cube if unspecified), timing out to sim-state targeting so # it never stalls. seen = bool(objects) and any( v.get("source") == "vision" and (want is None or v.get("color") == want) for v in objects.values()) waited = time.monotonic() - self._step_since if not seen and waited < 2.5: if self._think_n % 10 == 0: log.info("[PERCEIVE] looking for %s cube via ego camera… (%.1fs)", want or "a", waited) return "wait" resolved = self._resolve_target(objects, want) if resolved is None: from models.rx1_ik import CUBE_FALLBACK, RED_CUBE_POS color = want or "red" cube_pos = CUBE_FALLBACK.get(color, RED_CUBE_POS).copy() self._weld_col = self._weld_bit_by_color.get(color, self._nu) name, src = f"{color}_cube", "fallback-constant" else: name, cube_pos, self._weld_col = resolved src = (objects.get(name, {}) or {}).get("source", "state") place_pos = self._place_pos_for_step(step, objects) pregrasp = cube_pos + np.array([0.0, 0.0, 0.12]) approach = None if self._planner is not None and hasattr(self._planner, "plan_approach"): approach = self._planner.plan_approach(hand_now, pregrasp, n=12) self._executor.start(cube_pos, eq_id=0, box_pos=place_pos, approach_path=approach) dist_init = float(np.linalg.norm(cube_pos - hand_now)) log.info( "[PLAN] step='%s' target=%s source=%s cube=%s place=%s " "dist=%.4f m approach_wp=%s weld_bit=%s method=jacobian_dls", step.get("label", "?"), name, src, np.round(cube_pos, 3), np.round(place_pos, 3), dist_init, 0 if approach is None else len(approach), self._weld_col, ) return "started" # ── main act loop ───────────────────────────────────────────────────────── def act(self, world_embedding: np.ndarray, task_embedding: np.ndarray, proprio: Optional[np.ndarray] = None, text: str = "", objects: Optional[dict] = None, grounding: Optional[dict] = None, **_) -> np.ndarray: import mujoco as _mj self._think_n += 1 nu = self._nu if text and text != self._task_text and not self._executor.running: self.set_task_text(text) q_cur = (proprio[:nu].copy() if (proprio is not None and len(proprio) >= nu) else np.zeros(nu, dtype=np.float32)) # current hand position (for planning the approach), via FK on q_cur d = _mj.MjData(self._ik.model) d.qpos[:nu] = q_cur _mj.mj_kinematics(self._ik.model, d) hand_now = d.site_xpos[self._ik._rsite].copy() # ── 1) advance a running pick→place ─────────────────────────────────── if self._executor.running: target = self._executor.step(q_cur) if self._think_n % 5 == 0: site_pos = hand_now tgt_pos = self._executor._cube dist = float(np.linalg.norm(tgt_pos - site_pos)) log.info( "[ACT] phase=%-10s site=%s target=%s dist=%.4f m " "weld=%s queue=%d method=jacobian_dls", self._executor.phase.name, np.round(site_pos, 3), np.round(tgt_pos, 3), dist, self._executor.weld_active, len(self._queue), ) return self._emit(target) # ── 2) expand a pending high-level plan once perception is available ── if self._plan is not None and not self._queue: waited = time.monotonic() - self._step_since have_vision = bool(objects) and any( v.get("source") == "vision" for v in objects.values()) if not have_vision and waited < 2.5 and bool(objects): return self._emit(self._ik.rest()) # let the camera lock on self._queue = self._expand_plan(objects) self._plan = None self._cur_step_id = None if not self._queue: self._intent = "rest" # ── 3) start the next queued step ───────────────────────────────────── if self._queue: front = self._queue[0] if id(front) != self._cur_step_id: # new front step → reset timer self._cur_step_id = id(front) self._step_since = time.monotonic() result = self._start_step(front, q_cur, hand_now, objects, grounding) if result == "wait": return self._emit(self._ik.rest()) self._queue.pop(0) if result == "started": return self._emit(self._executor.step(q_cur)) return self._emit(self._ik.rest()) # skipped — try next cycle # ── 4) idle → hold the requested pose ───────────────────────────────── target = self._pose_for_intent(q_cur) if self._think_n % 10 == 0: log.info("[ACT] intent=%-14s method=pose_lookup think_n=%d", self._intent, self._think_n) return self._emit(target) def _emit(self, target: np.ndarray) -> np.ndarray: """Build the (chunk_len, action_dim) command, raising the active weld bit for the cube currently grasped.""" chunk = np.zeros((self._chunk_len, self._action_dim), dtype=np.float32) chunk[:, :self._nu] = target if self._executor.weld_active and self._weld_col is not None: chunk[:, self._weld_col] = 1.0 return chunk def _pose_for_intent(self, q_cur: np.ndarray) -> np.ndarray: if self._intent == "wave": # Raise the arm, then oscillate the forearm so it actually waves. # Amplitude/frequency are kept under the env slew limit (3 rad/s) so # the motion stays smooth: 0.3 rad * 2π * 1.0 Hz ≈ 1.9 rad/s. import time q = self._ik.wave() phase = 2.0 * np.pi * 1.0 * time.monotonic() idx = self._ik._idx if "act_right_forearm2forearmrot_joint" in idx: q[idx["act_right_forearm2forearmrot_joint"]] += 0.30 * np.sin(phase) if "act_right_forearmrot2forearm_pitch_joint" in idx: q[idx["act_right_forearmrot2forearm_pitch_joint"]] += 0.15 * np.sin(phase + 0.8) return q if self._intent == "raise_arms": return self._ik.raise_arms() if self._intent == "look_down": return self._ik.look_down(self._ik.rest()) return self._ik.rest() # ─────────────────────────────────────────────────────────────── factory class CosmosPolicy: def __new__(cls, config: dict) -> BasePolicy: # type: ignore[misc] if config.get("stub", True): if config.get("use_rx1_ik", False): log.info("Policy: RX1 IK (VLA-driven, all 40 joints dynamic, nu+2=42 actions)") return _RX1IKPolicy( chunk_len=config.get("action_chunk_len", 10), action_dim=config.get("action_dim", 42), ) log.info("Policy: STUB mode (random noise)") rng = np.random.default_rng(7) return _StubPolicy( chunk_len=config.get("action_chunk_len", 10), action_dim=config.get("action_dim", 21), noise=config.get("action_noise_stub", 0.01), rng=rng, ) return _RealPolicy( model_id=config["model_id"], device=config.get("device", "cuda"), dtype_str=config.get("dtype", "bfloat16"), cfg=config, )