File size: 7,691 Bytes
d8290eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """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/<ROBOT>/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,
)
|