from __future__ import annotations
import copy
import dataclasses
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Literal
import mujoco
import numpy as np
import torch
from mjlab.managers import CommandTerm
from mjlab.tasks.tracking.mdp import MotionCommandCfg as MjlabMotionCommandCfg
from mjlab.utils.lab_api.math import (
matrix_from_quat,
quat_apply,
quat_apply_inverse,
quat_error_magnitude,
quat_from_euler_xyz,
quat_inv,
quat_mul,
sample_uniform,
yaw_quat,
)
from mjlab.viewer.debug_visualizer import DebugVisualizer
from wbc_mjlab.env.mdp.sampling import (
RsiCfg,
TrackingSimilarityState,
bin_index_for_frame,
build_bin_valid_mask,
compile_similarity_terms,
compute_assist_gain_matrix,
compute_sampling_prob_matrix,
resolve_tracking_reward_indices,
rsi_failure_signal_label,
sample_adaptive_bins,
step_tracking_reward_similarity,
step_tracking_similarity,
trajectory_conditional_prob_row,
update_failure_ema,
)
from wbc_mjlab.motion.library import MotionFrame, OnlineMotion
from wbc_mjlab.viewer.motion_vis import (
clip_name_for_trajectory,
error_to_rgba,
format_motion_context_html,
format_rsi_panel_html,
)
if TYPE_CHECKING:
from collections.abc import Callable
from typing import Any
import viser
from mjlab.entity import Entity
from mjlab.envs import ManagerBasedRlEnv
_DESIRED_FRAME_COLORS = ((1.0, 0.5, 0.5), (0.5, 1.0, 0.5), (0.5, 0.5, 1.0))
[docs]
class MotionLoader:
"""Converted motion library: a bundled NPZ or a dataset directory of ``npz/*.npz`` clips.
The NPZ path of the command, kept for converted datasets and bundles; a dataset with source
clips is played online (:class:`wbc_mjlab.motion.library.OnlineMotion`).
"""
def __init__(
self,
motion_file: str,
body_indexes: torch.Tensor,
anchor_body_index: int,
step_dt: float,
device: str = "cpu",
) -> None:
from wbc_mjlab.motion.stack_bundle import (
list_clip_npz_files,
stack_clip_arrays_from_paths,
)
body_idx = body_indexes.detach().cpu().numpy()
source = Path(motion_file).expanduser().resolve()
self.motion_source = str(source)
if source.is_dir():
clip_paths = list_clip_npz_files(source)
if not clip_paths:
raise FileNotFoundError(
f"No clip NPZs in {source / 'npz'}; convert clips or pass a .npz bundle."
)
stacked = stack_clip_arrays_from_paths(clip_paths)
seg_start = stacked.segment_start_idx
seg_len = stacked.segment_length
self._load_arrays(
joint_pos=stacked.joint_pos,
joint_vel=stacked.joint_vel,
body_pos_w=stacked.body_pos_w,
body_quat_w=stacked.body_quat_w,
body_lin_vel_w=stacked.body_lin_vel_w,
body_ang_vel_w=stacked.body_ang_vel_w,
seg_start=seg_start,
seg_len=seg_len,
body_idx=body_idx,
anchor_body_index=anchor_body_index,
step_dt=step_dt,
device=device,
segment_names=stacked.clip_names,
)
return
if not source.is_file():
raise FileNotFoundError(f"Motion source not found: {source}")
with np.load(source, allow_pickle=True) as data:
if "segment_start_idx" in data and "segment_length" in data:
seg_start = np.asarray(data["segment_start_idx"], dtype=np.int64)
seg_len = np.asarray(data["segment_length"], dtype=np.int64)
else:
total = int(data["joint_pos"].shape[0])
seg_start = np.asarray([0], dtype=np.int64)
seg_len = np.asarray([total], dtype=np.int64)
if "segment_names" in data:
segment_names = tuple(str(x) for x in data["segment_names"].tolist())
elif "segment_source" in data:
from wbc_mjlab.motion.manifest import clip_name_from_path
segment_names = tuple(
clip_name_from_path(str(x)) for x in data["segment_source"].tolist()
)
else:
segment_names = (source.stem,)
self._load_arrays(
joint_pos=np.asarray(data["joint_pos"]),
joint_vel=np.asarray(data["joint_vel"]),
body_pos_w=np.asarray(data["body_pos_w"]),
body_quat_w=np.asarray(data["body_quat_w"]),
body_lin_vel_w=np.asarray(data["body_lin_vel_w"]),
body_ang_vel_w=np.asarray(data["body_ang_vel_w"]),
seg_start=seg_start,
seg_len=seg_len,
body_idx=body_idx,
anchor_body_index=anchor_body_index,
step_dt=step_dt,
device=device,
segment_names=segment_names,
)
def _load_arrays(
self,
*,
joint_pos: np.ndarray,
joint_vel: np.ndarray,
body_pos_w: np.ndarray,
body_quat_w: np.ndarray,
body_lin_vel_w: np.ndarray,
body_ang_vel_w: np.ndarray,
seg_start: np.ndarray,
seg_len: np.ndarray,
body_idx: np.ndarray,
anchor_body_index: int,
step_dt: float,
device: str,
segment_names: tuple[str, ...] = (),
) -> None:
self.joint_pos = torch.as_tensor(joint_pos, dtype=torch.float32, device=device)
self.joint_vel = torch.as_tensor(joint_vel, dtype=torch.float32, device=device)
self.body_pos_w = torch.as_tensor(
np.asarray(body_pos_w[:, body_idx], dtype=np.float32), device=device
)
self.body_quat_w = torch.as_tensor(
np.asarray(body_quat_w[:, body_idx], dtype=np.float32), device=device
)
self.body_lin_vel_w = torch.as_tensor(
np.asarray(body_lin_vel_w[:, body_idx], dtype=np.float32), device=device
)
self.body_ang_vel_w = torch.as_tensor(
np.asarray(body_ang_vel_w[:, body_idx], dtype=np.float32), device=device
)
self.time_step_total = int(self.joint_pos.shape[0])
self.segment_start_idx = torch.tensor(seg_start, dtype=torch.long, device=device)
self.segment_length = torch.tensor(seg_len, dtype=torch.long, device=device)
self.num_trajectories = int(self.segment_start_idx.shape[0])
self.segment_end_idx = self.segment_start_idx + self.segment_length
if len(segment_names) < self.num_trajectories:
segment_names = segment_names + tuple(
f"traj_{i}" for i in range(len(segment_names), self.num_trajectories)
)
self.segment_names = segment_names
anchor_lin_vel = self.body_lin_vel_w[:, anchor_body_index]
anchor_ang_vel = self.body_ang_vel_w[:, anchor_body_index]
self.anchor_lin_acc_w = torch.gradient(anchor_lin_vel, spacing=step_dt, dim=0)[0]
self.anchor_ang_acc_w = torch.gradient(anchor_ang_vel, spacing=step_dt, dim=0)[0]
[docs]
def frame(self, steps: torch.Tensor) -> MotionFrame:
return MotionFrame(
joint_pos=self.joint_pos[steps],
joint_vel=self.joint_vel[steps],
body_pos_w=self.body_pos_w[steps],
body_quat_w=self.body_quat_w[steps],
body_lin_vel_w=self.body_lin_vel_w[steps],
body_ang_vel_w=self.body_ang_vel_w[steps],
anchor_lin_acc_w=self.anchor_lin_acc_w[steps],
anchor_ang_acc_w=self.anchor_ang_acc_w[steps],
)
def _uses_online_motion(cfg: MotionCommandCfg) -> bool:
from wbc_mjlab.motion.library import has_sources
source = Path(cfg.motion_file).expanduser()
if cfg.motion_source == "npz":
return False
online = source.suffix.lower() != ".npz" and (
has_sources(source) or (source.is_file() and source.suffix.lower() != ".npz")
)
if cfg.motion_source == "online" and not online:
raise FileNotFoundError(
f"motion_source='online' needs source clips (CSV / PKL), none in {source}"
)
return online
def _robot_id(cfg: MotionCommandCfg) -> str:
if cfg.robot_id:
return cfg.robot_id
from wbc_mjlab.tasks import last_registered_robot_id
return last_registered_robot_id()
[docs]
class MotionCommand(CommandTerm):
"""Multi-clip motion playback with RSI resampling and assistive-wrench state.
Loads an NPZ motion library, streams reference kinematics each step, and
resamples start frames at episode boundaries according to :class:`RsiCfg`.
Observation and reward terms read reference state from this command via
``command_name="motion"``.
"""
cfg: MotionCommandCfg
_env: ManagerBasedRlEnv
def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
self.robot: Entity = env.scene[cfg.entity_name]
if cfg.actuated_joint_names:
tracked_joint_ids, _ = self.robot.find_joints(
cfg.actuated_joint_names, preserve_order=True
)
self._tracked_joint_ids = torch.tensor(
tracked_joint_ids, dtype=torch.long, device=self.device
)
else:
self._tracked_joint_ids = None
self.robot_anchor_body_index = self.robot.body_names.index(
self.cfg.anchor_body_name
)
self.motion_anchor_body_index = self.cfg.body_names.index(self.cfg.anchor_body_name)
self.body_indexes = torch.tensor(
self.robot.find_bodies(self.cfg.body_names, preserve_order=True)[0],
dtype=torch.long,
device=self.device,
)
self.motion: MotionLoader | OnlineMotion
if _uses_online_motion(cfg):
self.motion = OnlineMotion(
cfg.motion_file,
mj_model=env.sim.mj_model,
joint_names=self.robot.joint_names,
body_names=cfg.body_names,
anchor_body=cfg.anchor_body_name,
root_body=self.robot.body_names[0],
step_dt=env.step_dt,
nworld=self.num_envs,
device=self.device,
robot_id=_robot_id(cfg),
entity_name=cfg.entity_name,
mirror=cfg.mirror,
debias_z=cfg.debias_z,
storage=cfg.motion_storage,
format_name=cfg.source_format,
input_fps=cfg.source_fps,
)
else:
if cfg.mirror or cfg.debias_z:
print(
"[WARN] mirror / debias_z apply to source clips only; "
f"{cfg.motion_file} is converted NPZ (set at conversion)"
)
self.motion = MotionLoader(
self.cfg.motion_file,
self.body_indexes,
anchor_body_index=self.motion_anchor_body_index,
step_dt=env.step_dt,
device=self.device,
)
# Reference at ``time_steps``, recomputed when they change (tensor version counter).
self._ref_key: tuple[torch.Tensor, int] | None = None
self._ref_frame: MotionFrame | None = None
self.time_steps = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
self.trajectory_ids = torch.zeros(
self.num_envs, dtype=torch.long, device=self.device
)
# Per-episode world offset added to reference positions (``ground_reference``).
self.reference_offset_w = torch.zeros(self.num_envs, 3, device=self.device)
self.body_pos_relative_w = torch.zeros(
self.num_envs, len(cfg.body_names), 3, device=self.device
)
self.body_quat_relative_w = torch.zeros(
self.num_envs, len(cfg.body_names), 4, device=self.device
)
self.body_quat_relative_w[:, :, 0] = 1.0
self.bin_width_frames = max(
1, int(math.ceil(cfg.rsi.bin_width_s / max(env.step_dt, 1.0e-6)))
)
segment_bins = torch.ceil(
self.motion.segment_length.float() / float(self.bin_width_frames)
).long()
self.bins_per_trajectory = max(1, int(segment_bins.max().item()))
self.bin_valid_mask = build_bin_valid_mask(
self.motion.segment_length,
bins_per_trajectory=self.bins_per_trajectory,
bin_width_frames=self.bin_width_frames,
min_bin_span_ratio=cfg.rsi.min_bin_span_ratio,
device=self.device,
)
self._valid_bin_indices = self.bin_valid_mask.nonzero(as_tuple=False)
if self._valid_bin_indices.numel() == 0:
raise ValueError(
"No valid RSI bins after applying min_bin_span_ratio; "
"lower min_bin_span_ratio or use shorter bin_width_s."
)
self.bin_failure_levels = torch.zeros(
self.motion.num_trajectories,
self.bins_per_trajectory,
dtype=torch.float,
device=self.device,
)
self.bin_visit_counts = torch.zeros(
self.motion.num_trajectories,
self.bins_per_trajectory,
dtype=torch.float,
device=self.device,
)
self._episode_similarity_sum = torch.zeros(self.num_envs, device=self.device)
self._episode_similarity_denom = torch.ones(self.num_envs, device=self.device)
self._episode_step_count = torch.zeros(
self.num_envs, dtype=torch.long, device=self.device
)
self._episode_start_step = torch.zeros(
self.num_envs, dtype=torch.long, device=self.device
)
self._episode_start_bin = torch.zeros(
self.num_envs, dtype=torch.long, device=self.device
)
self.episode_assist_gain = torch.zeros(self.num_envs, device=self.device)
self.assist_force_w = torch.zeros(self.num_envs, 3, device=self.device)
self.assist_torque_w = torch.zeros(self.num_envs, 3, device=self.device)
self.metrics["error_anchor_pos"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_anchor_rot"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_anchor_lin_vel"] = torch.zeros(
self.num_envs, device=self.device
)
self.metrics["error_anchor_ang_vel"] = torch.zeros(
self.num_envs, device=self.device
)
self.metrics["error_body_pos"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_body_rot"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_joint_pos"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_joint_vel"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["sampling_entropy"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["sampling_top1_prob"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["sampling_top1_bin"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["assist_gain_mean"] = torch.zeros(self.num_envs, device=self.device)
self._similarity_terms, similarity_weight_sum = compile_similarity_terms(
cfg.rsi.similarity_terms
)
self._similarity_term_weight_sum = max(similarity_weight_sum, 1.0e-6)
if (
cfg.rsi.strategy == "similarity_ema"
and not cfg.rsi.similarity_from_rewards
and similarity_weight_sum <= 0.0
):
raise ValueError(
"similarity_ema requires at least one rsi.similarity_terms entry with weight > 0 "
"when similarity_from_rewards=False."
)
self._tracking_reward_indices: list[int] | None = None
self._tracking_reward_weight_sum = 0.0
self._ghost_model: mujoco.MjModel | None = None
self._ghost_color = np.array(cfg.viz.ghost_color, dtype=np.float32)
self._viewer_task_id: str | None = None
self._viewer_align_xy_yaw = False
self._viewer_color_bodies = False
self._viewer_context_html = None
self._viewer_rsi_html = None
@property
def bin_count(self) -> int:
return self.motion.num_trajectories * self.bins_per_trajectory
def _set_episode_similarity_denom(
self, env_ids: torch.Tensor, traj_ids: torch.Tensor, time_steps: torch.Tensor
) -> None:
if not self.cfg.rsi.similarity_norm_by_remaining_clip:
return
seg_start = self.motion.segment_start_idx[traj_ids]
seg_len = self.motion.segment_length[traj_ids]
local_step = torch.clamp(time_steps - seg_start, min=0)
remaining = seg_len - local_step
max_episode = float(self._env.max_episode_length)
norm = torch.minimum(
remaining.float(),
torch.full_like(remaining.float(), max_episode),
)
self._episode_similarity_denom[env_ids] = torch.clamp(norm, min=1.0)
def _update_failure_levels(self, env_ids: torch.Tensor) -> None:
if self.cfg.rsi.sampling_mode != "adaptive" or env_ids.numel() == 0:
return
active_mask = self._episode_step_count[env_ids] > 0
if not torch.any(active_mask):
return
env_ids = env_ids[active_mask]
rsi = self.cfg.rsi
strategy = rsi.strategy
update_failure_ema(
self.bin_failure_levels,
strategy=strategy,
bins_per_trajectory=self.bins_per_trajectory,
alpha=rsi.alpha,
traj_ids=self.trajectory_ids[env_ids],
start_bins=self._episode_start_bin[env_ids],
episode_terminated=(
self._env.termination_manager.terminated[env_ids]
if strategy == "binary_failure"
else None
),
episode_similarity_sum=(
self._episode_similarity_sum[env_ids] if strategy == "similarity_ema" else None
),
episode_step_count=(
self._episode_step_count[env_ids] if strategy == "similarity_ema" else None
),
similarity_denom=(
self._episode_similarity_denom[env_ids]
if rsi.similarity_norm_by_remaining_clip and strategy == "similarity_ema"
else None
),
norm_by_remaining_clip=(
rsi.similarity_norm_by_remaining_clip and strategy == "similarity_ema"
),
)
def _set_episode_assist_gain(
self, env_ids: torch.Tensor, traj_ids: torch.Tensor, bins: torch.Tensor
) -> None:
if not self.cfg.assistive_wrench_enabled:
self.episode_assist_gain[env_ids] = 0.0
return
failure = self.bin_failure_levels[traj_ids, bins]
similarity = 1.0 - failure
eta = max(self.cfg.assistive_eta, 1.0e-6)
self.episode_assist_gain[env_ids] = torch.clamp(
1.0 - similarity / eta, 0.0, self.cfg.assistive_beta_max
)
def _adaptive_sampling(self, env_ids: torch.Tensor):
self._update_failure_levels(env_ids)
rsi = self.cfg.rsi
traj_ids, bins, time_steps, probs_valid = sample_adaptive_bins(
self.bin_failure_levels,
self._valid_bin_indices,
segment_length=self.motion.segment_length,
segment_start_idx=self.motion.segment_start_idx,
bin_width_frames=self.bin_width_frames,
temperature_base=rsi.temperature_base,
uniform_ratio=rsi.uniform_ratio,
num_samples=len(env_ids),
device=self.device,
)
self.trajectory_ids[env_ids] = traj_ids
self.time_steps[env_ids] = time_steps
self._episode_start_bin[env_ids] = bins
self._set_episode_similarity_denom(env_ids, traj_ids, time_steps)
self._set_episode_assist_gain(env_ids, traj_ids, bins)
self._record_bin_visits(traj_ids, bins)
num_valid = max(1, probs_valid.shape[0])
H = -(probs_valid * (probs_valid + 1e-12).log()).sum()
self.metrics["sampling_entropy"][:] = (
float(H / math.log(num_valid)) if num_valid > 1 else 1.0
)
pmax, imax = probs_valid.max(dim=0)
self.metrics["sampling_top1_prob"][:] = float(pmax)
self.metrics["sampling_top1_bin"][:] = float(imax) / float(num_valid)
def _uniform_sampling(self, env_ids: torch.Tensor):
traj_ids = torch.randint(
0, self.motion.num_trajectories, (len(env_ids),), device=self.device
)
seg_lengths = self.motion.segment_length[traj_ids]
local_frames = (
sample_uniform(0.0, 1.0, (len(env_ids),), device=self.device)
* (seg_lengths.float() - 1.0)
).long()
self.trajectory_ids[env_ids] = traj_ids
self.time_steps[env_ids] = self.motion.segment_start_idx[traj_ids] + local_frames
start_bins = self._bin_index_for_frame(traj_ids, self.time_steps[env_ids])
self._episode_start_bin[env_ids] = start_bins
self._set_episode_similarity_denom(env_ids, traj_ids, self.time_steps[env_ids])
self._set_episode_assist_gain(env_ids, traj_ids, start_bins)
self._record_bin_visits(traj_ids, start_bins)
num_valid = max(1, int(self.bin_valid_mask.sum().item()))
self.metrics["sampling_entropy"][:] = 1.0
self.metrics["sampling_top1_prob"][:] = 1.0 / num_valid
self.metrics["sampling_top1_bin"][:] = 0.5
def _tracking_similarity_state(self) -> TrackingSimilarityState:
return TrackingSimilarityState(
tracked_joint_pos_error=self._tracked_joint_pos_error(),
anchor_pos_error=self.anchor_pos_w - self.robot_anchor_pos_w,
anchor_ori_error=quat_error_magnitude(
self.anchor_quat_w, self.robot_anchor_quat_w
),
body_pos_error=torch.sum(
torch.square(self.body_pos_relative_w - self.robot_body_pos_w), dim=-1
),
body_ori_error=(
quat_error_magnitude(self.body_quat_relative_w, self.robot_body_quat_w) ** 2
),
body_lin_vel_error=torch.sum(
torch.square(self.body_lin_vel_w - self.robot_body_lin_vel_w), dim=-1
),
body_ang_vel_error=torch.sum(
torch.square(self.body_ang_vel_w - self.robot_body_ang_vel_w), dim=-1
),
)
def _ensure_tracking_reward_indices(self) -> None:
if self._tracking_reward_indices is not None:
return
rsi = self.cfg.rsi
self._tracking_reward_indices, self._tracking_reward_weight_sum = (
resolve_tracking_reward_indices(
self._env.reward_manager,
name_prefix=rsi.tracking_reward_prefix,
)
)
if not self._tracking_reward_indices:
raise ValueError(
"similarity_from_rewards requires at least one reward term "
f"matching prefix {rsi.tracking_reward_prefix!r} with weight > 0."
)
def _step_tracking_similarity(self) -> torch.Tensor:
if self.cfg.rsi.similarity_from_rewards:
self._ensure_tracking_reward_indices()
return step_tracking_reward_similarity(
self._env.reward_manager._step_reward,
self._tracking_reward_indices,
self._tracking_reward_weight_sum,
)
return step_tracking_similarity(
self._similarity_terms,
self._similarity_term_weight_sum,
self._tracking_similarity_state(),
num_envs=self.num_envs,
device=self.device,
)
def _bin_index_for_frame(
self, trajectory_ids: torch.Tensor, time_steps: torch.Tensor
):
return bin_index_for_frame(
segment_start_idx=self.motion.segment_start_idx,
time_steps=time_steps,
trajectory_ids=trajectory_ids,
bin_width_frames=self.bin_width_frames,
bins_per_trajectory=self.bins_per_trajectory,
)
# --- WBC reference features (Table S3); stacked in ``command`` for the actor. ---
# SE layouts use ``ref_anchor_pos_w`` / ``ref_anchor_ori_6d`` via ``presets/se_actor.py`` instead.
@property
def ref_base_height(self) -> torch.Tensor:
"""Anchor height relative to env origin (z_I r̂_IB)."""
return self.anchor_pos_w[:, 2:3] - self._env.scene.env_origins[:, 2:3]
@property
def ref_base_lin_vel_b(self) -> torch.Tensor:
"""Reference anchor linear velocity in anchor frame (B v̂_IB)."""
return quat_apply_inverse(self.anchor_quat_w, self.anchor_lin_vel_w)
@property
def ref_base_ang_vel_b(self) -> torch.Tensor:
"""Reference anchor angular velocity in anchor frame (B ω̂_IB)."""
return quat_apply_inverse(self.anchor_quat_w, self.anchor_ang_vel_w)
@property
def ref_gravity_b(self) -> torch.Tensor:
"""Reference gravity in anchor frame (B ĝ_I)."""
return quat_apply_inverse(self.anchor_quat_w, self.robot.data.gravity_vec_w)
@property
def ref_base_lin_acc_b(self) -> torch.Tensor:
return quat_apply_inverse(self.anchor_quat_w, self.anchor_lin_acc_w)
@property
def ref_base_ang_acc_b(self) -> torch.Tensor:
return quat_apply_inverse(self.anchor_quat_w, self.anchor_ang_acc_w)
@property
def tracked_joint_pos(self) -> torch.Tensor:
"""Reference joint positions for actuated / tracked DoFs (absolute)."""
if self._tracked_joint_ids is not None:
return self.joint_pos[:, self._tracked_joint_ids]
return self.joint_pos
@property
def tracked_joint_vel(self) -> torch.Tensor:
"""Reference joint velocities for actuated / tracked DoFs (absolute)."""
if self._tracked_joint_ids is not None:
return self.joint_vel[:, self._tracked_joint_ids]
return self.joint_vel
@property
def command(self) -> torch.Tensor:
"""Legacy stacked WBC reference vector (same layout as default ref obs terms).
Prefer configuring individual reference observation terms in ``wbc_env_cfg``.
Kept for ONNX metadata, ``wbc_command_dim``, and legacy deploy bundles.
"""
return torch.cat(
[
self.ref_base_height,
self.ref_base_lin_vel_b,
self.ref_base_ang_vel_b,
self.ref_gravity_b,
self.tracked_joint_pos,
],
dim=-1,
)
@property
def wbc_command_dim(self) -> int:
return int(self.command.shape[-1])
[docs]
def reference(self) -> MotionFrame:
"""Reference at the current ``time_steps``: computed once per change of the frame
indices (any write to ``time_steps`` bumps its version), then shared by every term."""
ts = self.time_steps
if (
self._ref_frame is None
or self._ref_key is None
or self._ref_key[0] is not ts
or self._ref_key[1] != ts._version
):
self._ref_frame = self.motion.frame(ts)
self._ref_key = (ts, ts._version)
return self._ref_frame
def _patch_reference(self, env_ids: torch.Tensor) -> None:
"""Recompute only ``env_ids`` after their frame indices were just written, when that one
write is the only change since the reference was computed (else the next
:meth:`reference` recomputes every env)."""
ts, key, ref = self.time_steps, self._ref_key, self._ref_frame
if ref is None or key is None or key[0] is not ts or ts._version != key[1] + 1:
return
part = self.motion.frame(ts[env_ids])
for f in dataclasses.fields(MotionFrame):
getattr(ref, f.name)[env_ids] = getattr(part, f.name)
self._ref_key = (ts, ts._version)
@property
def joint_pos(self) -> torch.Tensor:
return self.reference().joint_pos
@property
def joint_vel(self) -> torch.Tensor:
return self.reference().joint_vel
def _tracked_joint_pos_error(self) -> torch.Tensor:
error = self.joint_pos - self.robot_joint_pos
if self._tracked_joint_ids is not None:
error = error[:, self._tracked_joint_ids]
return error
def _tracked_joint_vel_error(self) -> torch.Tensor:
error = self.joint_vel - self.robot_joint_vel
if self._tracked_joint_ids is not None:
error = error[:, self._tracked_joint_ids]
return error
@property
def body_pos_w(self) -> torch.Tensor:
return (
self.reference().body_pos_w
+ self._env.scene.env_origins[:, None, :]
+ self.reference_offset_w[:, None, :]
)
@property
def body_quat_w(self) -> torch.Tensor:
return self.reference().body_quat_w
@property
def body_lin_vel_w(self) -> torch.Tensor:
return self.reference().body_lin_vel_w
@property
def body_ang_vel_w(self) -> torch.Tensor:
return self.reference().body_ang_vel_w
@property
def anchor_pos_w(self) -> torch.Tensor:
return (
self.reference().body_pos_w[:, self.motion_anchor_body_index]
+ self._env.scene.env_origins
+ self.reference_offset_w
)
@property
def anchor_quat_w(self) -> torch.Tensor:
return self.reference().body_quat_w[:, self.motion_anchor_body_index]
@property
def anchor_lin_vel_w(self) -> torch.Tensor:
return self.reference().body_lin_vel_w[:, self.motion_anchor_body_index]
@property
def anchor_ang_vel_w(self) -> torch.Tensor:
return self.reference().body_ang_vel_w[:, self.motion_anchor_body_index]
@property
def anchor_lin_acc_w(self) -> torch.Tensor:
return self.reference().anchor_lin_acc_w
@property
def anchor_ang_acc_w(self) -> torch.Tensor:
return self.reference().anchor_ang_acc_w
@property
def robot_joint_pos(self) -> torch.Tensor:
return self.robot.data.joint_pos
@property
def robot_joint_vel(self) -> torch.Tensor:
return self.robot.data.joint_vel
@property
def robot_body_pos_w(self) -> torch.Tensor:
return self.robot.data.body_link_pos_w[:, self.body_indexes]
@property
def robot_body_quat_w(self) -> torch.Tensor:
return self.robot.data.body_link_quat_w[:, self.body_indexes]
@property
def robot_body_lin_vel_w(self) -> torch.Tensor:
return self.robot.data.body_link_lin_vel_w[:, self.body_indexes]
@property
def robot_body_ang_vel_w(self) -> torch.Tensor:
return self.robot.data.body_link_ang_vel_w[:, self.body_indexes]
@property
def robot_anchor_pos_w(self) -> torch.Tensor:
return self.robot.data.body_link_pos_w[:, self.robot_anchor_body_index]
@property
def robot_anchor_quat_w(self) -> torch.Tensor:
return self.robot.data.body_link_quat_w[:, self.robot_anchor_body_index]
@property
def robot_anchor_lin_vel_w(self) -> torch.Tensor:
return self.robot.data.body_link_lin_vel_w[:, self.robot_anchor_body_index]
@property
def robot_anchor_ang_vel_w(self) -> torch.Tensor:
return self.robot.data.body_link_ang_vel_w[:, self.robot_anchor_body_index]
def _update_metrics(self):
self.metrics["error_anchor_pos"] = torch.norm(
self.anchor_pos_w - self.robot_anchor_pos_w, dim=-1
)
self.metrics["error_anchor_rot"] = quat_error_magnitude(
self.anchor_quat_w, self.robot_anchor_quat_w
)
self.metrics["error_anchor_lin_vel"] = torch.norm(
self.anchor_lin_vel_w - self.robot_anchor_lin_vel_w, dim=-1
)
self.metrics["error_anchor_ang_vel"] = torch.norm(
self.anchor_ang_vel_w - self.robot_anchor_ang_vel_w, dim=-1
)
self.metrics["error_body_pos"] = torch.norm(
self.body_pos_relative_w - self.robot_body_pos_w, dim=-1
).mean(dim=-1)
self.metrics["error_body_rot"] = quat_error_magnitude(
self.body_quat_relative_w, self.robot_body_quat_w
).mean(dim=-1)
self.metrics["error_body_lin_vel"] = torch.norm(
self.body_lin_vel_w - self.robot_body_lin_vel_w, dim=-1
).mean(dim=-1)
self.metrics["error_body_ang_vel"] = torch.norm(
self.body_ang_vel_w - self.robot_body_ang_vel_w, dim=-1
).mean(dim=-1)
self.metrics["error_joint_pos"] = torch.norm(
self._tracked_joint_pos_error(), dim=-1
)
self.metrics["error_joint_vel"] = torch.norm(
self._tracked_joint_vel_error(), dim=-1
)
self.metrics["assist_gain_mean"] = self.episode_assist_gain
def _write_reference_state_to_sim(
self,
env_ids: torch.Tensor,
root_pos: torch.Tensor,
root_ori: torch.Tensor,
root_lin_vel: torch.Tensor,
root_ang_vel: torch.Tensor,
joint_pos: torch.Tensor,
joint_vel: torch.Tensor,
) -> None:
soft_limits = self.robot.data.soft_joint_pos_limits[env_ids]
joint_pos = torch.clip(joint_pos, soft_limits[:, :, 0], soft_limits[:, :, 1])
self.robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
root_state = torch.cat([root_pos, root_ori, root_lin_vel, root_ang_vel], dim=-1)
self.robot.write_root_state_to_sim(root_state, env_ids=env_ids)
self.robot.reset(env_ids=env_ids)
def _resample_command(self, env_ids: torch.Tensor):
rsi = self.cfg.rsi
if rsi.sampling_mode == "start":
self.trajectory_ids[env_ids] = 0
self.time_steps[env_ids] = self.motion.segment_start_idx[0]
self._episode_start_bin[env_ids] = 0
start_traj = self.trajectory_ids[env_ids]
start_bins = self._episode_start_bin[env_ids]
self._set_episode_similarity_denom(env_ids, start_traj, self.time_steps[env_ids])
self._set_episode_assist_gain(env_ids, start_traj, start_bins)
self._record_bin_visits(start_traj, start_bins)
elif rsi.sampling_mode == "uniform":
self._uniform_sampling(env_ids)
else:
assert rsi.sampling_mode == "adaptive"
self._adaptive_sampling(env_ids)
self._patch_reference(env_ids)
self._update_reference_offset(env_ids)
root_pos = self.body_pos_w[env_ids, 0].clone()
root_ori = self.body_quat_w[env_ids, 0].clone()
root_lin_vel = self.body_lin_vel_w[env_ids, 0].clone()
root_ang_vel = self.body_ang_vel_w[env_ids, 0].clone()
range_list = [
self.cfg.pose_range.get(key, (0.0, 0.0))
for key in ["x", "y", "z", "roll", "pitch", "yaw"]
]
ranges = torch.tensor(range_list, device=self.device)
rand_samples = sample_uniform(
ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=self.device
)
root_pos += rand_samples[:, 0:3]
orientations_delta = quat_from_euler_xyz(
rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5]
)
root_ori = quat_mul(orientations_delta, root_ori)
range_list = [
self.cfg.velocity_range.get(key, (0.0, 0.0))
for key in ["x", "y", "z", "roll", "pitch", "yaw"]
]
ranges = torch.tensor(range_list, device=self.device)
rand_samples = sample_uniform(
ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=self.device
)
root_lin_vel += rand_samples[:, :3]
root_ang_vel += rand_samples[:, 3:]
joint_pos = self.joint_pos[env_ids].clone()
joint_vel = self.joint_vel[env_ids]
joint_pos += sample_uniform(
lower=self.cfg.joint_position_range[0],
upper=self.cfg.joint_position_range[1],
size=joint_pos.shape,
device=joint_pos.device,
)
self._write_reference_state_to_sim(
env_ids,
root_pos,
root_ori,
root_lin_vel,
root_ang_vel,
joint_pos,
joint_vel,
)
self._episode_start_step[env_ids] = self.time_steps[env_ids]
self._episode_similarity_sum[env_ids] = 0.0
self._episode_step_count[env_ids] = 0
self.update_relative_body_poses()
self._seed_default_relative_action(env_ids)
def _update_reference_offset(self, env_ids: torch.Tensor) -> None:
"""Center the start-frame anchor XY on the env origin (``ground_reference``).
The offset is held for the whole episode, so the reference keeps its own
displacement but starts on the sub-terrain origin instead of wherever the
clip was recorded.
"""
if not self.cfg.ground_reference:
return
start_xy = self.reference().body_pos_w[env_ids, self.motion_anchor_body_index, :2]
self.reference_offset_w[env_ids] = 0.0
self.reference_offset_w[env_ids, :2] = -start_xy
def _seed_default_relative_action(self, env_ids: torch.Tensor) -> None:
"""Warm-start default-relative actions after RSI (no-op for residual actions)."""
action_manager = getattr(self._env, "action_manager", None)
if action_manager is None or "joint_pos" not in action_manager.active_terms:
return
seed = getattr(
action_manager.get_term("joint_pos"), "seed_from_motion_reference", None
)
if seed is not None:
seed(env_ids)
[docs]
def update_relative_body_poses(self) -> None:
anchor_pos_w_repeat = self.anchor_pos_w[:, None, :].repeat(
1, len(self.cfg.body_names), 1
)
anchor_quat_w_repeat = self.anchor_quat_w[:, None, :].repeat(
1, len(self.cfg.body_names), 1
)
robot_anchor_pos_w_repeat = self.robot_anchor_pos_w[:, None, :].repeat(
1, len(self.cfg.body_names), 1
)
robot_anchor_quat_w_repeat = self.robot_anchor_quat_w[:, None, :].repeat(
1, len(self.cfg.body_names), 1
)
delta_pos_w = robot_anchor_pos_w_repeat
delta_pos_w[..., 2] = anchor_pos_w_repeat[..., 2]
delta_ori_w = yaw_quat(
quat_mul(robot_anchor_quat_w_repeat, quat_inv(anchor_quat_w_repeat))
)
self.body_quat_relative_w = quat_mul(delta_ori_w, self.body_quat_w)
self.body_pos_relative_w = delta_pos_w + quat_apply(
delta_ori_w, self.body_pos_w - anchor_pos_w_repeat
)
[docs]
def compute(
self, dt: float | torch.Tensor, env_ids: torch.Tensor | None = None
) -> None:
# Stash dt so _update_command can skip motion-frame advances on dt=0
# (explicit reset and auto-reset envs). Base compute owns timers.
self._last_compute_dt = dt
super().compute(dt, env_ids)
def _update_command(self, env_ids: torch.Tensor | None) -> None:
"""Advance clip frames on the per-step path; refresh relative poses always.
``env_ids`` is None on the regular step (all envs) and the reset ids on
``reset()``. Motion time only advances where ``dt > 0`` so a reset does
not skip the frame the robot was just written to.
"""
dt = getattr(self, "_last_compute_dt", 0.0)
if env_ids is not None:
advance = None
elif isinstance(dt, torch.Tensor):
advance = dt > 0.0
elif dt > 0.0:
advance = torch.ones(self.num_envs, dtype=torch.bool, device=self.device)
else:
advance = None
if advance is not None and bool(advance.any()):
if self.cfg.rsi.strategy == "similarity_ema":
sim = self._step_tracking_similarity()
self._episode_similarity_sum[advance] += sim[advance]
self._episode_step_count[advance] += 1
self.time_steps[advance] += 1
seg_end = self.motion.segment_end_idx[self.trajectory_ids]
wrap_ids = torch.where(advance & (self.time_steps >= seg_end))[0]
if wrap_ids.numel() > 0:
self._resample_command(wrap_ids)
self.update_relative_body_poses()
[docs]
def set_viewer_task_id(self, task_id: str | None) -> None:
self._viewer_task_id = task_id
def _record_bin_visits(self, traj_ids: torch.Tensor, bins: torch.Tensor) -> None:
for traj_id, bin_id in zip(traj_ids.tolist(), bins.tolist(), strict=True):
self.bin_visit_counts[int(traj_id), int(bin_id)] += 1.0
def _rsi_view_tensors(
self,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
rsi = self.cfg.rsi
failure = self.bin_failure_levels
prob = compute_sampling_prob_matrix(
failure,
self.bin_valid_mask,
temperature_base=rsi.temperature_base,
uniform_ratio=rsi.uniform_ratio,
)
assist = compute_assist_gain_matrix(
failure,
eta=self.cfg.assistive_eta,
beta_max=self.cfg.assistive_beta_max,
enabled=self.cfg.assistive_wrench_enabled,
)
return (
failure.detach().cpu(),
prob.detach().cpu(),
self.bin_visit_counts.detach().cpu(),
assist.detach().cpu(),
self.bin_valid_mask.detach().cpu(),
)
def _clip_name(self, traj_id: int) -> str:
return clip_name_for_trajectory(self.motion.segment_names, traj_id)
def _segment_phase(self, env_idx: int) -> float:
traj_id = int(self.trajectory_ids[env_idx].item())
seg_start = int(self.motion.segment_start_idx[traj_id].item())
seg_len = max(int(self.motion.segment_length[traj_id].item()) - 1, 1)
local = int(self.time_steps[env_idx].item()) - seg_start
return float(np.clip(local / seg_len, 0.0, 1.0))
def _body_tracking_errors(self, env_idx: int) -> tuple[np.ndarray, np.ndarray]:
pos_err = torch.norm(
self.body_pos_relative_w[env_idx] - self.robot_body_pos_w[env_idx], dim=-1
)
rot_err = quat_error_magnitude(
self.body_quat_relative_w[env_idx], self.robot_body_quat_w[env_idx]
)
return pos_err.cpu().numpy(), rot_err.cpu().numpy()
[docs]
def update_viewer_gui(self, env_idx: int) -> None:
if self._viewer_context_html is None:
return
traj_id = int(self.trajectory_ids[env_idx].item())
frame = int(self.time_steps[env_idx].item())
self._viewer_context_html.content = format_motion_context_html(
env_idx=env_idx,
traj_id=traj_id,
clip_name=self._clip_name(traj_id),
frame=frame,
phase=self._segment_phase(env_idx),
task_id=self._viewer_task_id,
rsi_mode=self.cfg.rsi.sampling_mode,
rsi_strategy=self.cfg.rsi.strategy,
anchor_body=self.cfg.anchor_body_name,
num_bodies=len(self.cfg.body_names),
)
if self._viewer_rsi_html is None:
return
bin_idx = int(
self._bin_index_for_frame(
self.trajectory_ids[env_idx : env_idx + 1],
self.time_steps[env_idx : env_idx + 1],
).item()
)
failure, prob, visits, assist, valid = self._rsi_view_tensors()
rsi = self.cfg.rsi
traj_probs = trajectory_conditional_prob_row(prob, valid, traj_id)
self._viewer_rsi_html.content = format_rsi_panel_html(
bin_idx=bin_idx,
num_bins=self.bins_per_trajectory,
bin_width_s=rsi.bin_width_s,
failure_levels=failure[traj_id].tolist(),
sampling_probs=traj_probs,
visit_counts=visits[traj_id].tolist(),
assist_gains=assist[traj_id].tolist(),
valid_mask=valid[traj_id].tolist(),
failure_signal_label=rsi_failure_signal_label(
rsi.strategy,
similarity_from_rewards=rsi.similarity_from_rewards,
),
show_assist=self.cfg.assistive_wrench_enabled,
beta_max=self.cfg.assistive_beta_max,
)
def _debug_vis_impl(self, visualizer: DebugVisualizer) -> None:
env_indices = visualizer.get_env_indices(self.num_envs)
if not env_indices:
return
if self.cfg.viz.mode == "ghost":
if self._ghost_model is None:
self._ghost_model = copy.deepcopy(self._env.sim.mj_model)
self._ghost_model.geom_rgba[:] = self._ghost_color
entity: Entity = self._env.scene[self.cfg.entity_name]
indexing = entity.indexing
free_joint_q_adr = indexing.free_joint_q_adr.cpu().numpy()
joint_q_adr = indexing.joint_q_adr.cpu().numpy()
for batch in env_indices:
if self._viewer_align_xy_yaw:
root_pos = self.body_pos_relative_w[batch, 0].cpu().numpy()
root_quat = self.body_quat_relative_w[batch, 0].cpu().numpy()
else:
root_pos = self.body_pos_w[batch, 0].cpu().numpy()
root_quat = self.body_quat_w[batch, 0].cpu().numpy()
qpos = np.zeros(self._env.sim.mj_model.nq)
qpos[free_joint_q_adr[0:3]] = root_pos
qpos[free_joint_q_adr[3:7]] = root_quat
qpos[joint_q_adr] = self.joint_pos[batch].cpu().numpy()
visualizer.add_ghost_mesh(qpos, model=self._ghost_model, label=f"ghost_{batch}")
if self._viewer_color_bodies:
self._add_body_error_markers(visualizer, batch)
elif self.cfg.viz.mode == "frames":
for batch in env_indices:
if self._viewer_align_xy_yaw:
desired_body_pos = self.body_pos_relative_w[batch].cpu().numpy()
desired_body_quat = self.body_quat_relative_w[batch]
else:
desired_body_pos = self.body_pos_w[batch].cpu().numpy()
desired_body_quat = self.body_quat_w[batch]
desired_body_rotm = matrix_from_quat(desired_body_quat).cpu().numpy()
current_body_pos = self.robot_body_pos_w[batch].cpu().numpy()
current_body_quat = self.robot_body_quat_w[batch]
current_body_rotm = matrix_from_quat(current_body_quat).cpu().numpy()
pos_err, rot_err = self._body_tracking_errors(batch)
for i, body_name in enumerate(self.cfg.body_names):
err = float(pos_err[i] + rot_err[i])
axis_colors = None
if self._viewer_color_bodies:
rgba = error_to_rgba(err)
rgb = (rgba[0], rgba[1], rgba[2])
axis_colors = (rgb, rgb, rgb)
visualizer.add_frame(
position=desired_body_pos[i],
rotation_matrix=desired_body_rotm[i],
scale=0.08,
label=f"desired_{body_name}_{batch}",
axis_colors=axis_colors or _DESIRED_FRAME_COLORS,
)
visualizer.add_frame(
position=current_body_pos[i],
rotation_matrix=current_body_rotm[i],
scale=0.12,
label=f"current_{body_name}_{batch}",
axis_colors=axis_colors,
)
if self._viewer_color_bodies:
visualizer.add_sphere(
center=current_body_pos[i],
radius=0.025,
color=error_to_rgba(err),
label=f"err_{body_name}_{batch}",
)
if self._viewer_align_xy_yaw:
desired_anchor_pos = self.body_pos_relative_w[
batch, self.motion_anchor_body_index
]
desired_anchor_quat = self.body_quat_relative_w[
batch, self.motion_anchor_body_index
]
else:
desired_anchor_pos = self.anchor_pos_w[batch]
desired_anchor_quat = self.anchor_quat_w[batch]
desired_rotation_matrix = matrix_from_quat(desired_anchor_quat).cpu().numpy()
visualizer.add_frame(
position=desired_anchor_pos.cpu().numpy(),
rotation_matrix=desired_rotation_matrix,
scale=0.1,
label=f"desired_anchor_{batch}",
axis_colors=_DESIRED_FRAME_COLORS,
)
current_anchor_pos = self.robot_anchor_pos_w[batch].cpu().numpy()
current_anchor_quat = self.robot_anchor_quat_w[batch]
current_rotation_matrix = matrix_from_quat(current_anchor_quat).cpu().numpy()
visualizer.add_frame(
position=current_anchor_pos,
rotation_matrix=current_rotation_matrix,
scale=0.15,
label=f"current_anchor_{batch}",
)
def _add_body_error_markers(self, visualizer: DebugVisualizer, env_idx: int) -> None:
pos_err, rot_err = self._body_tracking_errors(env_idx)
body_pos = self.robot_body_pos_w[env_idx].cpu().numpy()
for i, body_name in enumerate(self.cfg.body_names):
err = float(pos_err[i] + rot_err[i])
visualizer.add_sphere(
center=body_pos[i],
radius=0.03,
color=error_to_rgba(err),
label=f"err_{body_name}_{env_idx}",
)
[docs]
def create_gui(
self,
name: str,
server: viser.ViserServer,
get_env_idx: Callable[[], int],
on_change: Callable[[], None] | None = None,
request_action: Callable[[str, Any], None] | None = None,
) -> None:
max_frame = int(self.motion.time_step_total) - 1
with server.gui.add_folder(name.capitalize()):
with server.gui.add_folder("Selected env", expand_by_default=True):
self._viewer_context_html = server.gui.add_html("")
with server.gui.add_folder("Adaptive sampling (RSI)", expand_by_default=False):
self._viewer_rsi_html = server.gui.add_html("")
align_cb = server.gui.add_checkbox(
"Align xy/yaw to reference",
initial_value=self._viewer_align_xy_yaw,
)
color_cb = server.gui.add_checkbox(
"Color bodies by tracking error",
initial_value=self._viewer_color_bodies,
)
@align_cb.on_update
def _on_align(_) -> None:
self._viewer_align_xy_yaw = bool(align_cb.value)
if on_change is not None:
on_change()
@color_cb.on_update
def _on_color(_) -> None:
self._viewer_color_bodies = bool(color_cb.value)
if on_change is not None:
on_change()
scrubber = server.gui.add_slider(
"Frame",
min=0,
max=max_frame,
step=1,
initial_value=0,
)
@scrubber.on_update
def _(_) -> None:
idx = get_env_idx()
self.time_steps[idx] = int(scrubber.value)
if on_change is not None:
on_change()
all_envs_cb = server.gui.add_checkbox("All envs", initial_value=True)
start_btn = server.gui.add_button("Start Here")
@start_btn.on_click
def _(_) -> None:
if request_action is not None:
request_action(
"CUSTOM",
{"type": "gui_reset", "all_envs": all_envs_cb.value},
)
self._scrubber_handles = (scrubber, all_envs_cb, start_btn)
self._set_scrubber_disabled(True)
self.update_viewer_gui(get_env_idx())
def _set_scrubber_disabled(self, disabled: bool) -> None:
for handle in self._scrubber_handles:
handle.disabled = disabled
[docs]
def on_viewer_pause(self, paused: bool) -> None:
if hasattr(self, "_scrubber_handles"):
self._set_scrubber_disabled(not paused)
[docs]
def apply_gui_reset(self, env_ids: torch.Tensor) -> bool:
if not hasattr(self, "_scrubber_handles"):
return False
frame = int(self._scrubber_handles[0].value)
self.reset_to_frame(env_ids, frame)
self.update_relative_body_poses()
return True
[docs]
def reset_to_frame(self, env_ids: torch.Tensor, frame: int) -> None:
self.time_steps[env_ids] = frame
self._patch_reference(env_ids)
traj_ids = torch.searchsorted(
self.motion.segment_end_idx,
torch.full((len(env_ids),), frame, device=self.device),
right=False,
)
traj_ids = torch.clamp(traj_ids, max=self.motion.num_trajectories - 1)
self.trajectory_ids[env_ids] = traj_ids
self._update_reference_offset(env_ids)
self._write_reference_state_to_sim(
env_ids,
self.body_pos_w[env_ids, 0],
self.body_quat_w[env_ids, 0],
self.body_lin_vel_w[env_ids, 0],
self.body_ang_vel_w[env_ids, 0],
self.joint_pos[env_ids],
self.joint_vel[env_ids],
)
[docs]
@dataclass(kw_only=True)
class MotionCommandCfg(MjlabMotionCommandCfg):
"""Config for :class:`MotionCommand` (set per robot in ``<robot>_base_cfg``)."""
motion_file: str
"""Motion library: a dataset folder of source clips (CSV / PKL, at its root or in
``raw/``) played online, or converted NPZ (a ``.npz`` bundle or a folder of ``npz/*.npz``)."""
motion_source: Literal["auto", "online", "npz"] = "auto"
"""``auto``: online when the folder has source clips, else its converted NPZ."""
mirror: bool = False
"""Online: add a left/right mirrored copy of every clip (virtual segments, no storage)."""
debias_z: bool = False
"""Online: shift each clip vertically so its lowest foot sits at the robot's sole height."""
motion_storage: Literal["device", "host"] = "device"
"""Online: source frames on the env device, or in pinned host memory (huge datasets)."""
source_format: str | None = None
"""Online: source format name (``motion_formats``); inferred from the extension if unset."""
source_fps: float | None = None
"""Online: override the source frame rate (default: the format's / file's)."""
robot_id: str = ""
"""Robot of the motion (symmetry for ``mirror``, feet for ``debias_z``); set per robot."""
anchor_body_name: str
"""Body used for anchor-frame errors and assistive wrench."""
body_names: tuple[str, ...]
"""Keybodies tracked in rewards, RSI, and critic observations."""
entity_name: str
"""Scene entity name for the robot (usually ``\"robot\"``)."""
actuated_joint_names: tuple[str, ...] = ()
"""If set, joint tracking metrics/RSI use only these DoFs (subset of the robot)."""
pose_range: dict[str, tuple[float, float]] = field(default_factory=dict)
"""Domain randomization ranges for reference root pose at resample."""
velocity_range: dict[str, tuple[float, float]] = field(default_factory=dict)
"""Domain randomization ranges for reference root velocity at resample."""
joint_position_range: tuple[float, float] = (-0.52, 0.52)
"""Domain randomization range for joint position offsets at resample."""
rsi: RsiCfg = field(default_factory=RsiCfg)
"""Reference-state initialization / adaptive bin sampling."""
assistive_wrench_enabled: bool = True
"""Whether assistive wrench curriculum is active."""
assistive_beta_max: float = 0.6
"""Max assistive gain β."""
assistive_eta: float = 0.8
"""Assistive wrench curriculum exponent."""
ground_reference: bool = False
"""Shift each episode's reference so its start-frame anchor XY sits on the env
origin. Needed on generated terrain, where env origins are sub-terrain centers
and clips recorded far from (0, 0) would otherwise start off-tile."""
[docs]
@dataclass
class VizCfg:
mode: Literal["ghost", "frames"] = "ghost"
ghost_color: tuple[float, float, float, float] = (0.45, 0.6, 0.9, 0.5)
viz: VizCfg = field(default_factory=VizCfg)
[docs]
def build(self, env: ManagerBasedRlEnv) -> MotionCommand:
return MotionCommand(self, env)
