"""Asynchronous real-ball adapter for the native H2 Online-CLoSD runtime.""" from __future__ import annotations from dataclasses import dataclass import random import threading import time import numpy as np from .embodiments import EmbodimentSpec from .mathutil import quat_normalize, quat_slerp_sign_aligned from .reference import FrozenPose, ReferenceClip, ReferenceGenerator @dataclass(frozen=True) class ExternalCommand: command_id: int received_s: float deadline_s: float hit_pos: np.ndarray hit_vel: np.ndarray landing_xy: np.ndarray base_y: float prefix_history: np.ndarray class RealReferenceState: """Own the 50 Hz policy reference while CLoSD generates off-thread. A planner packet's time-to-hit remains a live wall-clock deadline. The CLoSD generator receives the initial condition, and any generation latency is compensated by skipping already elapsed reference frames before a clip is activated. """ def __init__( self, spec: EmbodimentSpec, generator: ReferenceGenerator, frozen: FrozenPose, *, landing_xy: tuple[float, float] = (2.87, 0.0), latency_skip_s: float = 0.0, generation_trigger_tth_s: float | None = None, ) -> None: self.spec = spec self.generator = generator self.frozen = frozen self.landing_xy = np.asarray(landing_xy, dtype=np.float64).reshape(2) self.latency_skip_s = max(0.0, float(latency_skip_s)) self.generation_trigger_tth_s = generation_trigger_tth_s self._notice_pending = False self._notice_command_id = 0 self._live_base_y = 0.0 self.dt = float(spec.control_dt) seed = generator.prefix_frame(frozen.pelvis_pos, frozen.pelvis_quat, frozen.joint_pos) self.history = np.repeat(seed[None], generator.history_len, axis=0).astype(np.float32) self.ref_joint_pos = frozen.joint_pos.astype(np.float64).copy() self.ref_joint_vel = np.zeros_like(self.ref_joint_pos) self.ref_pelvis_pos = frozen.pelvis_pos.astype(np.float64).copy() self.ref_pelvis_quat = quat_normalize(frozen.pelvis_quat.astype(np.float64)) self._clip: ReferenceClip | None = None self._frame = 0 self._active_elapsed_s = 0.0 self._active_command_id = 0 self._active_deadline_s = float("-inf") self._active_hit_pos = np.zeros(3, dtype=np.float64) self._active_hit_vel = np.zeros(3, dtype=np.float64) self._active_landing_xy = self.landing_xy.copy() self._last_requested_command_id = 0 self._live_command_id = 0 self._live_deadline_s = float("-inf") self._live_hit_pos = np.zeros(3, dtype=np.float64) self._live_hit_vel = np.zeros(3, dtype=np.float64) self._live_landing_xy = self.landing_xy.copy() self._frozen = True self._freeze_elapsed_s = float(spec.frozen_blend_time_s) self._freeze_start_joint_pos = self.ref_joint_pos.copy() self._freeze_start_joint_vel = self.ref_joint_vel.copy() self._freeze_start_pelvis_pos = self.ref_pelvis_pos.copy() self._freeze_start_pelvis_quat = self.ref_pelvis_quat.copy() self._frozen_root_xy = self.ref_pelvis_pos[:2].copy() self._last_robot_root_xy: np.ndarray | None = None self._lock = threading.Lock() self._generation_thread: threading.Thread | None = None self._queued: ExternalCommand | None = None self._pending: tuple[ExternalCommand, ReferenceClip, float] | None = None self._generation_error: BaseException | None = None self._generation_epoch = 0 self._closing = False self.last_generation_s = 0.0 self.last_skip_frames = 0 self.last_wait_frames = 0 self.last_motion_hit_time_s = 0.0 self.last_remaining_command_s = 0.0 self.last_clamped = False self.last_clamp_delta_m = 0.0 self.timing_trace = None @property def command_id(self) -> int: return int(self._active_command_id) @property def is_frozen(self) -> bool: return bool(self._frozen) @property def active_hit_deadline_s(self) -> float: """Monotonic deadline actually used by the active H2 reference.""" return float(self._active_deadline_s) @property def active_landing_xy(self) -> np.ndarray: """Planner target paired with the active reference, in robot court metres.""" return self._active_landing_xy.copy() def has_generation_in_flight(self) -> bool: with self._lock: return bool( (self._generation_thread is not None and self._generation_thread.is_alive()) or self._queued is not None or self._pending is not None ) def raise_if_failed(self) -> None: with self._lock: error = self._generation_error self._generation_error = None if error is not None: raise RuntimeError("asynchronous H2 CLoSD generation failed") from error def freeze_for_mocap_loss(self) -> None: """Cancel strike admission and blend to the existing frozen reference. The policy keeps tracking this reference. An in-flight generator is allowed to finish, but its result/error belongs to the old epoch and cannot reactivate a strike. Keep that worker alive so a later command never runs concurrently against the same generator. """ self._notice_pending = False self._notice_command_id = 0 self._live_command_id = 0 self._live_deadline_s = float("-inf") with self._lock: self._generation_epoch += 1 self._queued = None self._pending = None self._generation_error = None self._begin_freeze() self._clip = None def warmup(self, passes: int = 1) -> None: passes = max(0, int(passes)) if passes == 0: return # A warmup is a latency/compilation operation, not a sample from the # deployed scenario. Preserve every global RNG used by the native # generator so enabling warmup cannot silently select a different # reference motion for the same configured seed. import torch python_rng_state = random.getstate() numpy_rng_state = np.random.get_state() torch_cpu_rng_state = torch.random.get_rng_state() torch_cuda_rng_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None center = np.array( [ 1.0, 0.5 * sum(self.spec.hit_time_range), self.spec.hit_plane_x, 0.0, 0.5 * sum(self.spec.hit_z_range), 0.0, 0.0, 0.0, 0.0, ], dtype=np.float32, ) started = time.perf_counter() try: for index in range(passes): self.generator.generate(center, self.history.copy(), self.generator.pred_len) print(f"[CLoSD WARMUP] {index + 1}/{passes} complete") finally: random.setstate(python_rng_state) np.random.set_state(numpy_rng_state) torch.random.set_rng_state(torch_cpu_rng_state) if torch_cuda_rng_states is not None: torch.cuda.set_rng_state_all(torch_cuda_rng_states) print(f"[CLoSD WARMUP] total={time.perf_counter() - started:.3f}s") def update_robot_state( self, root_xy: np.ndarray, pelvis_quat_wxyz: np.ndarray, joint_pos_policy: np.ndarray, ) -> None: root_xy = np.asarray(root_xy, dtype=np.float64).reshape(2) quat = quat_normalize(np.asarray(pelvis_quat_wxyz, dtype=np.float64).reshape(4)) joints = np.asarray(joint_pos_policy, dtype=np.float64).reshape(self.generator.num_joints) if not np.isfinite(root_xy).all() or not np.isfinite(quat).all() or not np.isfinite(joints).all(): raise RuntimeError("cannot add non-finite robot state to the CLoSD prefix") self._last_robot_root_xy = root_xy.copy() if self._frozen: # A result can activate on this tick: blend out of the latest held # reference, not the previous tick's measured pelvis position. self._refresh() pelvis = np.array([root_xy[0], root_xy[1], self.ref_pelvis_pos[2]], dtype=np.float64) pelvis[2] = float(pelvis[2]) if np.isfinite(pelvis[2]) else float(self.spec.nominal_root_z) frame = self.generator.prefix_frame(pelvis, quat, joints) if self.history.shape[0] > 1: self.history[:-1] = self.history[1:] self.history[-1] = frame def update_external_command( self, command_id: int, time_to_hit_s: float, hit_pos: np.ndarray, hit_vel: np.ndarray, *, base_y: float, landing_xy: np.ndarray | None = None, ) -> None: command_id = int(command_id) tth = float(time_to_hit_s) pos = np.asarray(hit_pos, dtype=np.float64).reshape(3) vel = np.asarray(hit_vel, dtype=np.float64).reshape(3) landing = ( self.landing_xy.copy() if landing_xy is None else np.asarray(landing_xy, dtype=np.float64).reshape(2) ) if command_id <= 0 or tth <= 0.0: raise ValueError("active CLoSD commands require command_id>0 and time_to_hit>0") if not np.isfinite(np.concatenate([pos, vel, landing, [tth, base_y]])).all(): raise ValueError("external command contains non-finite values") now = time.monotonic() with self._lock: if command_id != self._live_command_id: # A new command supersedes old results, including errors that # completed just before this control tick. Same-ID refinements # must keep their in-flight generation and completed result. self._generation_epoch += 1 self._queued = self._pending = None self._generation_error = None self._live_command_id = command_id self._live_deadline_s = now + tth self._live_hit_pos = pos.copy() self._live_hit_vel = vel.copy() self._live_landing_xy = landing.copy() self._live_base_y = float(base_y) if command_id == self._active_command_id: self._active_deadline_s = self._live_deadline_s self._active_hit_pos = pos.copy() self._active_hit_vel = vel.copy() self._active_landing_xy = landing.copy() if command_id == self._last_requested_command_id: return if self.generation_trigger_tth_s is not None: first_notice = not self._notice_pending or command_id != self._notice_command_id self._notice_pending = True self._notice_command_id = command_id self._begin_freeze() if first_notice: print(f"[CLoSD NOTICE] cmd_id={command_id} tth={tth:.3f}s " f"generate_at={self.generation_trigger_tth_s:.3f}s") if self.timing_trace is not None: self.timing_trace.record("command_notice", reference_id=command_id, time_to_hit_s=tth) # step() runs after this tick's measured state enters the prefix. return self._request_generation(now) def _request_generation(self, now: float) -> None: request = ExternalCommand( command_id=self._live_command_id, received_s=now, deadline_s=self._live_deadline_s, hit_pos=self._live_hit_pos.copy(), hit_vel=self._live_hit_vel.copy(), landing_xy=self._live_landing_xy.copy(), base_y=self._live_base_y, prefix_history=self.history.copy(), ) self._last_requested_command_id = request.command_id if self.generation_trigger_tth_s is not None: print(f"[CLoSD GENERATE] cmd_id={request.command_id} tth={request.deadline_s-now:.3f}s") if self.timing_trace is not None: self.timing_trace.record("generation_request", reference_id=request.command_id, time_to_hit_s=request.deadline_s-now) with self._lock: if self._closing: return self._queued = request if self._generation_thread is None or not self._generation_thread.is_alive(): self._generation_thread = threading.Thread( target=self._generation_worker, name="h2-closd-generator", daemon=True, ) self._generation_thread.start() def _generation_worker(self) -> None: while True: with self._lock: if self._closing or self._queued is None: self._generation_thread = None return request = self._queued self._queued = None epoch = self._generation_epoch try: requested_tth = max(0.05, request.deadline_s - request.received_s) command = np.array( [1.0, requested_tth, *request.hit_pos, 0.0, 0.0, 0.0, 0.0], dtype=np.float32, ) started = time.perf_counter() if self.timing_trace is not None: self.timing_trace.record("generation_start", reference_id=request.command_id, time_to_hit_s=request.deadline_s-time.monotonic()) clip = self.generator.generate( command, request.prefix_history, self.generator.pred_len, ) elapsed = time.perf_counter() - started if self.timing_trace is not None: self.timing_trace.record("generation_end", reference_id=request.command_id) except BaseException as error: # surfaced synchronously on the 50 Hz loop if self.timing_trace is not None: self.timing_trace.record("generation_error", reference_id=request.command_id) with self._lock: if epoch != self._generation_epoch or request.command_id != self._live_command_id: continue self._generation_error = error self._generation_thread = None return with self._lock: if epoch != self._generation_epoch or request.command_id != self._live_command_id: continue # Keep only the newest completed request. A newer queued command # must never be overwritten by an older diffusion result. if self._queued is None or request.command_id >= self._queued.command_id: self._pending = (request, clip, elapsed) def _consume_pending(self, now: float) -> bool: with self._lock: pending = self._pending self._pending = None if pending is None: return False request, clip, elapsed = pending if request.command_id != self._live_command_id: return False self._validate_clip(clip) motion_hit_s = float(np.asarray(clip.used_condition, dtype=np.float64).reshape(9)[1]) remaining_command_s = float(self._live_deadline_s - now - self.latency_skip_s) if not np.isfinite([motion_hit_s, remaining_command_s]).all() or motion_hit_s <= 0.0: raise RuntimeError("H2 CLoSD returned an invalid effective hit time") if self.generation_trigger_tth_s is not None and remaining_command_s > motion_hit_s: # A revised prediction can postpone contact after generation. Keep # following the measured root, not a world-fixed prefix of holds. with self._lock: if self._pending is None: self._pending = pending return False # Align the generator's *effective* (possibly clamped) hit time to the # predictor's live deadline. A positive phase delta starts inside the # clip; a negative delta holds the current reference before playback. # Always discard generated frame 0, matching training/G1 external # deployment. An extra hold frame preserves exact phase when the # rounded delta is zero. phase_frames = int(round((motion_hit_s - remaining_command_s) * self.generator.fps)) if phase_frames >= 1: skipped = phase_frames wait_frames = 0 else: skipped = 1 wait_frames = 1 - phase_frames # A playable reference requires two frames after phase alignment. # Preserve valid wait prefixes, but drop a lone late motion frame. if skipped >= clip.length or clip.length - skipped + wait_frames < 2: self._notice_pending = False print( f"[CLoSD DROP] cmd_id={request.command_id} result expired; " f"motion_hit={motion_hit_s:.3f}s remaining={remaining_command_s:.3f}s " f"skip={skipped}/{clip.length}" ) return False # Training applies the exit blend in original clip time. Bake it before # latency trimming so a two-frame skip starts at the two-frame alpha, # instead of silently restarting a fresh 240 ms blend. self._blend_clip_start(clip) if skipped > 0: clip = ReferenceClip( joint_pos=clip.joint_pos[skipped:].copy(), joint_vel=clip.joint_vel[skipped:].copy(), pelvis_pos=clip.pelvis_pos[skipped:].copy(), pelvis_quat=clip.pelvis_quat[skipped:].copy(), clamped=clip.clamped, clamp_delta=clip.clamp_delta, raw_condition=clip.raw_condition.copy(), used_condition=clip.used_condition.copy(), ) if wait_frames > 0: clip = ReferenceClip( joint_pos=np.concatenate( [np.repeat(self.ref_joint_pos[None], wait_frames, axis=0), clip.joint_pos], axis=0 ), joint_vel=np.concatenate( [np.repeat(self.ref_joint_vel[None], wait_frames, axis=0), clip.joint_vel], axis=0 ), pelvis_pos=np.concatenate( [np.repeat(self.ref_pelvis_pos[None], wait_frames, axis=0), clip.pelvis_pos], axis=0 ), pelvis_quat=np.concatenate( [np.repeat(self.ref_pelvis_quat[None], wait_frames, axis=0), clip.pelvis_quat], axis=0 ), clamped=clip.clamped, clamp_delta=clip.clamp_delta, raw_condition=clip.raw_condition.copy(), used_condition=clip.used_condition.copy(), ) self._validate_clip(clip) self._clip = clip self._frame = 0 self._active_elapsed_s = 0.0 self._active_command_id = request.command_id self._active_deadline_s = self._live_deadline_s self._active_hit_pos = self._live_hit_pos.copy() self._active_hit_vel = self._live_hit_vel.copy() self._active_landing_xy = self._live_landing_xy.copy() self._frozen = False self._notice_pending = False self._freeze_elapsed_s = 0.0 self.last_generation_s = float(elapsed) self.last_skip_frames = int(skipped) self.last_wait_frames = int(wait_frames) self.last_motion_hit_time_s = motion_hit_s self.last_remaining_command_s = remaining_command_s self.last_clamped = bool(clip.clamped) self.last_clamp_delta_m = float(clip.clamp_delta) if self.timing_trace is not None: self.timing_trace.record("reference_adopt", reference_id=request.command_id, skip_frames=skipped, wait_frames=wait_frames) print( f"[CLoSD READY] cmd_id={request.command_id} gen={elapsed:.3f}s " f"motion_hit={motion_hit_s:.3f}s remaining={remaining_command_s:.3f}s " f"skip={skipped} wait={wait_frames} frames={clip.length} clamped={int(clip.clamped)} " f"delta={clip.clamp_delta:.3f}m" ) return True def _validate_clip(self, clip: ReferenceClip) -> None: """Reject malformed generator output before it can become a policy reference.""" joint_pos = np.asarray(clip.joint_pos) joint_vel = np.asarray(clip.joint_vel) pelvis_pos = np.asarray(clip.pelvis_pos) pelvis_quat = np.asarray(clip.pelvis_quat) raw_condition = np.asarray(clip.raw_condition) used_condition = np.asarray(clip.used_condition) if joint_pos.ndim != 2: raise RuntimeError(f"H2 CLoSD joint_pos shape is invalid: {joint_pos.shape}") length = int(joint_pos.shape[0]) expected = { "joint_pos": (length, self.generator.num_joints), "joint_vel": (length, self.generator.num_joints), "pelvis_pos": (length, 3), "pelvis_quat": (length, 4), "raw_condition": (9,), "used_condition": (9,), } actual = { "joint_pos": joint_pos.shape, "joint_vel": joint_vel.shape, "pelvis_pos": pelvis_pos.shape, "pelvis_quat": pelvis_quat.shape, "raw_condition": raw_condition.shape, "used_condition": used_condition.shape, } for name, expected_shape in expected.items(): if actual[name] != expected_shape: raise RuntimeError( f"H2 CLoSD {name} shape is invalid: {actual[name]} != {expected_shape}" ) if length < 2: raise RuntimeError("H2 CLoSD clip must contain at least two frames") arrays = (joint_pos, joint_vel, pelvis_pos, pelvis_quat, raw_condition, used_condition) if not all(np.isfinite(array).all() for array in arrays): raise RuntimeError("H2 CLoSD clip contains non-finite values") quaternion_norms = np.linalg.norm(pelvis_quat, axis=1) if np.any(quaternion_norms < 0.5) or np.any(quaternion_norms > 1.5): raise RuntimeError("H2 CLoSD pelvis quaternion norm is outside [0.5,1.5]") clamp_delta = float(clip.clamp_delta) if not np.isfinite(clamp_delta) or clamp_delta < 0.0: raise RuntimeError("H2 CLoSD clamp_delta is invalid") def _blend_clip_start(self, clip: ReferenceClip) -> None: duration = max(0.0, float(self.spec.frozen_exit_blend_time_s)) if duration <= 0.0 or clip.length == 0: return start_jp = self.ref_joint_pos.copy() start_jv = self.ref_joint_vel.copy() start_pp = self.ref_pelvis_pos.copy() start_pq = self.ref_pelvis_quat.copy() for index in range(clip.length): alpha = min(1.0, (index / self.generator.fps) / duration) clip.joint_pos[index] = start_jp + alpha * (clip.joint_pos[index] - start_jp) clip.joint_vel[index] = start_jv + alpha * (clip.joint_vel[index] - start_jv) clip.pelvis_pos[index] = start_pp + alpha * (clip.pelvis_pos[index] - start_pp) clip.pelvis_quat[index] = quat_slerp_sign_aligned(start_pq, clip.pelvis_quat[index], alpha) def _begin_freeze(self) -> None: if self._frozen: return self._frozen = True self._freeze_elapsed_s = 0.0 self._freeze_start_joint_pos = self.ref_joint_pos.copy() self._freeze_start_joint_vel = self.ref_joint_vel.copy() self._freeze_start_pelvis_pos = self.ref_pelvis_pos.copy() self._freeze_start_pelvis_quat = self.ref_pelvis_quat.copy() self._frozen_root_xy = self.ref_pelvis_pos[:2].copy() def _refresh(self) -> None: if not self._frozen and self._clip is not None: index = int(np.clip(self._frame, 0, self._clip.length - 1)) self.ref_joint_pos = self._clip.joint_pos[index].astype(np.float64).copy() self.ref_joint_vel = self._clip.joint_vel[index].astype(np.float64).copy() self.ref_pelvis_pos = self._clip.pelvis_pos[index].astype(np.float64).copy() self.ref_pelvis_quat = quat_normalize(self._clip.pelvis_quat[index].astype(np.float64)) return target_pelvis = self.frozen.pelvis_pos.astype(np.float64).copy() # Match Mycc StrokeController's frozen_reference_follows_robot: keep # the 0.24 s entry blend, but re-anchor its target XY on every sample. # Active clips above remain world-anchored; height/yaw/joints still # blend to the existing ready pose. The fallback is for pre-sensor use. target_pelvis[:2] = ( self._frozen_root_xy if self._last_robot_root_xy is None else self._last_robot_root_xy ) duration = max(0.0, float(self.spec.frozen_blend_time_s)) alpha = 1.0 if duration <= 0.0 else min(1.0, self._freeze_elapsed_s / duration) self.ref_joint_pos = self._freeze_start_joint_pos + alpha * ( self.frozen.joint_pos - self._freeze_start_joint_pos ) self.ref_joint_vel = (1.0 - alpha) * self._freeze_start_joint_vel self.ref_pelvis_pos = self._freeze_start_pelvis_pos + alpha * ( target_pelvis - self._freeze_start_pelvis_pos ) self.ref_pelvis_quat = quat_slerp_sign_aligned( self._freeze_start_pelvis_quat, self.frozen.pelvis_quat, alpha, ) def step(self, dt: float | None = None) -> None: self.raise_if_failed() now = time.monotonic() if self._notice_pending and self._live_command_id != self._last_requested_command_id: remaining = self._live_deadline_s - now if remaining <= 0.0: self._notice_pending = False elif remaining <= self.generation_trigger_tth_s: self._request_generation(now) activated = self._consume_pending(now) step_dt = self.dt if dt is None else max(0.0, float(dt)) entered_freeze = False if not self._frozen and self._clip is not None and not activated: self._active_elapsed_s += step_dt self._frame = int(round(self._active_elapsed_s * self.generator.fps)) if self._frame >= self._clip.length: self._frame = self._clip.length - 1 self._refresh() self._begin_freeze() entered_freeze = True elif self._active_deadline_s - now < -0.59: self._refresh() self._begin_freeze() entered_freeze = True if self._frozen and not entered_freeze: self._freeze_elapsed_s += step_dt self._refresh() def task_command(self) -> np.ndarray: now = time.monotonic() if self._notice_pending: result = np.concatenate( [[max(-0.59, self._live_deadline_s-now)], self._live_hit_pos, self._live_hit_vel, self._live_landing_xy] ) elif not self._frozen and self._active_command_id > 0: tth = max(-0.59, self._active_deadline_s - now) result = np.concatenate( [[tth], self._active_hit_pos, self._active_hit_vel, self._active_landing_xy] ) else: result = np.asarray(self.spec.frozen_command(), dtype=np.float64) if result.shape != (9,) or not np.isfinite(result).all(): raise RuntimeError("H2 task command is not a finite 9-vector") return result.astype(np.float32) def close(self) -> None: self._notice_pending = False with self._lock: self._closing = True self._queued = None thread = self._generation_thread if thread is not None and thread.is_alive(): thread.join(timeout=2.0) __all__ = ["ExternalCommand", "RealReferenceState"]