"""Online CLoSD reference generation and the frozen ready pose. ``ReferenceGenerator`` wraps the *training* runtime (``MotionBricksCLoSDRuntime`` with a native ``*.closd-runtime.pt`` bundle), so conditioning, clamping, sampling, decoding and FK are byte-for-byte the code used during RL. ``FrozenPose`` loads the per-robot ready pose that the training environment blends into whenever a generated clip is exhausted. """ from __future__ import annotations from dataclasses import dataclass import numpy as np import torch from .embodiments import EmbodimentSpec from .paths import load_closd_runtime_module @dataclass class ReferenceClip: joint_pos: np.ndarray # (T, N) policy joint order, absolute angles joint_vel: np.ndarray # (T, N) pelvis_pos: np.ndarray # (T, 3) env-local (= world for a single env) pelvis_quat: np.ndarray # (T, 4) wxyz clamped: bool # condition was clamped to the generator's domain clamp_delta: float # norm of the clamped position part [m] raw_condition: np.ndarray # (9,) used_condition: np.ndarray # (9,) @property def length(self) -> int: return int(self.joint_pos.shape[0]) @dataclass class FrozenPose: joint_pos: np.ndarray # (N,) policy joint order pelvis_pos: np.ndarray # (3,) env-local pelvis_quat: np.ndarray # (4,) wxyz def load_frozen_pose(spec: EmbodimentSpec, policy_joint_names: list[str], frame: int = 0) -> FrozenPose: """Load ``data//frozen_closd_ready_pose_toeout_lowload.npz`` (frame 0, like training).""" data = np.load(spec.ready_pose_path(), allow_pickle=True) joint_pos = np.asarray(data["joint_pos"], dtype=np.float64)[frame] body_pos = np.asarray(data["body_pos_w"], dtype=np.float64)[frame] body_quat = np.asarray(data["body_quat_w"], dtype=np.float64)[frame] if "joint_names" in data.files: file_joints = [str(x) for x in np.asarray(data["joint_names"]).ravel()] index = [file_joints.index(name) for name in policy_joint_names if name in file_joints] if len(index) != len(policy_joint_names): missing = [n for n in policy_joint_names if n not in file_joints] if missing and len(index) + len(missing) == len(policy_joint_names): # joints absent from the file (none expected) keep 0.0 full = np.zeros(len(policy_joint_names)) for i, name in enumerate(policy_joint_names): if name in file_joints: full[i] = joint_pos[file_joints.index(name)] joint_pos = full else: raise ValueError(f"ready pose is missing joints {missing}") else: joint_pos = joint_pos[index] elif joint_pos.shape[0] != len(policy_joint_names): raise ValueError( f"ready pose has {joint_pos.shape[0]} joints without names; policy expects {len(policy_joint_names)}" ) if "body_names" in data.files: body_names = [str(x) for x in np.asarray(data["body_names"]).ravel()] pelvis = body_names.index("pelvis") if "pelvis" in body_names else 0 else: pelvis = 0 # legacy G1 file: MuJoCo body order, pelvis first return FrozenPose(joint_pos=joint_pos, pelvis_pos=body_pos[pelvis].copy(), pelvis_quat=body_quat[pelvis].copy()) class ReferenceGenerator: """Generate reference clips with the training runtime and the native bundle.""" def __init__( self, spec: EmbodimentSpec, policy_joint_names: list[str], default_joint_pos: np.ndarray, device: str | None = None, sampling_steps: int | None = None, ): self.spec = spec self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") runtime_module = load_closd_runtime_module() self.runtime = runtime_module.MotionBricksCLoSDRuntime( str(spec.bundle_path()), humanoid_xml=str(spec.mjcf_path()), device=self.device, sampling_steps=sampling_steps, fps=spec.closd_fps, condition_clamp_enabled=True, condition_time_range=spec.condition_time_range, condition_lateral_range=spec.condition_lateral_range, condition_forward_range=spec.condition_forward_range, condition_height_range=spec.condition_height_range, robot_joint_names=list(policy_joint_names), robot_default_joint_pos=np.asarray(default_joint_pos, dtype=np.float32), reference_body_names=list(spec.reference_body_names), ) self.policy_joint_names = list(policy_joint_names) self.num_joints = len(policy_joint_names) # indices (into the policy joint vector) of the joints the model generates, in model order self.model_joint_index = np.asarray(self.runtime.isaaclab_to_mujoco, dtype=np.int64) self.context_len = int(getattr(self.runtime, "context_len", 0)) self.history_len = max(1, self.context_len) self.pred_len = int(self.runtime.pred_len) self.fps = float(self.runtime.fps) self.pelvis_index = list(spec.reference_body_names).index("pelvis") self.qpos_dim = 7 + len(self.model_joint_index) def prefix_frame(self, pelvis_pos_local: np.ndarray, pelvis_quat: np.ndarray, joint_pos_policy: np.ndarray) -> np.ndarray: """One prefix qpos frame: ``[root xyz, root wxyz, model joints]`` (training ``_current_robot_qpos_local``).""" return np.concatenate( [np.asarray(pelvis_pos_local, dtype=np.float32), np.asarray(pelvis_quat, dtype=np.float32), np.asarray(joint_pos_policy, dtype=np.float32)[self.model_joint_index]] ) @torch.no_grad() def generate(self, command_local: np.ndarray, prefix_history: np.ndarray, played_frames: int) -> ReferenceClip: """Generate one clip. ``command_local``: 9-vector ``[1, hit_time, hit_x, hit_y, hit_z, 0, 0, 0, 0]`` (env-local). ``prefix_history``: ``(history_len, 7 + n_model_joints)`` robot qpos frames. ``played_frames``: frames that will actually be played (``round(next_resample_time * fps)``). """ played = int(np.clip(played_frames, 1, self.pred_len)) cmd = torch.as_tensor(np.asarray(command_local, dtype=np.float32).reshape(1, 9), device=self.device) prefix = torch.as_tensor(np.asarray(prefix_history, dtype=np.float32).reshape(1, self.history_len, -1), device=self.device) played_t = torch.tensor([played], dtype=torch.long, device=self.device) payload = self.runtime.generate_batch_tensors(cmd, prefix_qpos=prefix, played_frames=played_t) joint_pos = payload["joint_pos"][0, :played].detach().cpu().numpy().astype(np.float64) joint_vel = payload["joint_vel"][0, :played].detach().cpu().numpy().astype(np.float64) body_pos = payload["body_pos_w"][0, :played].detach().cpu().numpy().astype(np.float64) body_quat = payload["body_quat_w"][0, :played].detach().cpu().numpy().astype(np.float64) raw = payload["_raw_condition"][0].detach().cpu().numpy().astype(np.float64) used = payload["_condition"][0].detach().cpu().numpy().astype(np.float64) delta = used[1:5] - raw[1:5] return ReferenceClip( joint_pos=joint_pos, joint_vel=joint_vel, pelvis_pos=body_pos[:, self.pelvis_index], pelvis_quat=body_quat[:, self.pelvis_index], clamped=bool(np.any(np.abs(delta) > 1e-6)), clamp_delta=float(np.linalg.norm(delta[1:4])), raw_condition=raw, used_condition=used, )