chore: publish H2 real-world RL system

d8290eb verified 3 days ago

7.69 kB

"""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,
        )

	"""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,
	)