feat(unitree_g1): standalone PICO SMPL publisher + dedup/replay fixes

Add a self-contained rt/smpl publisher in the pico_headset teleoperator
(pico_publisher.py + numpy SMPL FK in smpl_fk.py + vendored skeleton table)
so headset whole-body teleop no longer depends on gear_sonic/torch; only
xrobotoolkit_sdk is needed at the headset.

Also: share lowstate_to_obs/get_gravity_orientation via g1_utils (dedup
sonic_pipeline and UnitreeG1.get_observation), and fix dataset-replay joint
ordering (Unitree -> IsaacLab) for sonic.py --replay-dataset.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Martino Russi
2026-07-09 19:06:37 +02:00
parent 3363688f1e
commit 943ae78cfe
13 changed files with 798 additions and 86 deletions
@@ -30,7 +30,10 @@ from enum import IntEnum
import numpy as np
import onnxruntime as ort
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
from lerobot.robots.unitree_g1.g1_utils import (
G1_29_JointIndex,
get_gravity_orientation,
)
# ── Constants ────────────────────────────────────────────────────────────────
@@ -174,12 +177,6 @@ def _to_mujoco(a):
return a[MUJOCO_TO_ISAACLAB]
def _to_runtime(a):
r = np.zeros(29, np.float32)
r[MUJOCO_TO_ISAACLAB] = a
return r
DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
@@ -214,28 +211,6 @@ def compute_kp_kd():
_kp_kd = compute_kp_kd # backward-compatible alias
def lowstate_to_obs(lowstate) -> dict:
"""Build a robot observation dict from Unitree lowstate."""
obs: dict = {}
for motor in G1_29_JointIndex:
idx = motor.value
obs[f"{motor.name}.q"] = float(lowstate.motor_state[idx].q)
obs[f"{motor.name}.dq"] = float(lowstate.motor_state[idx].dq)
quat = lowstate.imu_state.quaternion
obs["imu.quat.w"] = float(quat[0])
obs["imu.quat.x"] = float(quat[1])
obs["imu.quat.y"] = float(quat[2])
obs["imu.quat.z"] = float(quat[3])
gyro = lowstate.imu_state.gyroscope
obs["imu.gyro.x"] = float(gyro[0])
obs["imu.gyro.y"] = float(gyro[1])
obs["imu.gyro.z"] = float(gyro[2])
wr = getattr(lowstate, "wireless_remote", None)
if wr is not None:
obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
return obs
# ── Quaternion helpers ────────────────────────────────────────────────────────
@@ -257,12 +232,6 @@ def quat_mul(q1, q2):
)
def gravity_dir(q):
q = q / (np.linalg.norm(q) + 1e-8)
qv = np.array([0, 0, 0, -1], dtype=np.float32)
return quat_mul(quat_mul(quat_conj(q), qv), q)[1:]
def quat_to_6d(q):
w, x, y, z = q
return np.array(
@@ -631,7 +600,7 @@ class StandingEncoderDecoder:
obs[off : off + sz] = h[f]
off += sz
for q in reversed(self.h_quat):
obs[off : off + 3] = gravity_dir(q)
obs[off : off + 3] = get_gravity_orientation(q)
off += 3
assert off == 994, f"Decoder obs mismatch: {off}"
return obs
@@ -689,7 +658,9 @@ class StandingEncoderDecoder:
print(
f" anchor_6d(identity): {self._anchor_6d(np.array([1, 0, 0, 0], np.float32), np.array([1, 0, 0, 0], np.float32))}"
)
print(f" gravity(identity): {gravity_dir(np.array([1, 0, 0, 0], np.float32))} (expect [0,0,-1])")
print(
f" gravity(identity): {get_gravity_orientation(np.array([1, 0, 0, 0], np.float32))} (expect [0,0,-1])"
)
dec0 = self.build_decoder_obs()
print(f" decoder q-delta max: {np.max(np.abs(dec0[94:384])):.6f}")
print(f" decoder dq max: {np.max(np.abs(dec0[384:674])):.6f}")
@@ -1122,7 +1093,6 @@ class PlannerController(StandingEncoderDecoder):
# ── Keyboard ──────────────────────────────────────────────────────────────────
class RawKeyboard:
def __init__(self):
self.fd = sys.stdin.fileno()
@@ -37,10 +37,10 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
_snapshot_ms,
clamp_mode_params,
compute_kp_kd,
lowstate_to_obs,
process_joystick,
should_replan_request,
)
from lerobot.robots.unitree_g1.g1_utils import lowstate_to_obs
logger = logging.getLogger(__name__)
@@ -185,25 +185,49 @@ class SonicWholeBodyController:
self._smpl_stream = SmplStream(host=host, port=port)
logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, port)
def _enter_wholebody(self) -> None:
"""Switch into SMPL whole-body tracking (encode_mode 2)."""
self.controller.encode_mode = 2
self.controller.reinit_heading = True
logger.info("SONIC: SMPL stream active -> whole-body tracking (mode 2)")
def _exit_wholebody(self) -> None:
"""Revert to locomotion/standing (encode_mode 0) after SMPL is lost.
Mirrors the 'M' toggle in sonic.py so the handoff is clean: the robot holds
a standing reference and (if a joystick teleop is attached) can be driven.
"""
self.controller.encode_mode = 0
self.controller.playing = True
self.controller.reinit_heading = True
self.ms.needs_replan = True
logger.warning("SONIC: SMPL stream lost/stale -> reverting to locomotion (standing)")
def run_step(self, action: dict, lowstate) -> dict:
if lowstate is None:
return {}
obs = lowstate_to_obs(lowstate)
# Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a
# direct rt/smpl subscription when SONIC_SMPL_STREAM is enabled.
# direct rt/smpl subscription when SONIC_SMPL_STREAM is enabled. A stale
# stream (headset silent past its timeout) is treated as "no SMPL" so the
# robot doesn't stay frozen tracking the last pose.
smpl = _extract_smpl_from_action(action)
if smpl is None and self._smpl_stream is not None:
window = self._smpl_stream.step()
if self._smpl_stream.has_data:
if self._smpl_stream.has_data and not self._smpl_stream.is_stale:
smpl = window
if smpl is not None:
# Full-body whole-body tracking: SMPL drives the reference, not joystick.
self.controller.encode_mode = 2
if self.controller.encode_mode != 2:
self._enter_wholebody()
self.controller.smpl_joints_10frame_step1 = smpl
return self._runtime.tick(obs, debug=False, use_joystick=False)
# No (or stale) SMPL: fall back to locomotion so the robot stays balanced.
if self.controller.encode_mode == 2:
self._exit_wholebody()
return self._runtime.tick(obs, debug=False)
def reset(self):
+41
View File
@@ -63,6 +63,47 @@ class G1_29_JointArmIndex(IntEnum):
kRightWristYaw = 28
def lowstate_to_obs(lowstate) -> dict:
"""Build a robot observation dict from a Unitree lowstate.
Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
lowstate -> obs mapping lives in exactly one place. Keys match the
``<joint>.q``/``imu.*`` schema consumed across the controllers.
"""
obs: dict = {}
for motor in G1_29_JointIndex:
idx = motor.value
obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
imu = lowstate.imu_state
if imu.gyroscope:
obs["imu.gyro.x"] = imu.gyroscope[0]
obs["imu.gyro.y"] = imu.gyroscope[1]
obs["imu.gyro.z"] = imu.gyroscope[2]
if imu.accelerometer:
obs["imu.accel.x"] = imu.accelerometer[0]
obs["imu.accel.y"] = imu.accelerometer[1]
obs["imu.accel.z"] = imu.accelerometer[2]
if imu.quaternion:
obs["imu.quat.w"] = imu.quaternion[0]
obs["imu.quat.x"] = imu.quaternion[1]
obs["imu.quat.y"] = imu.quaternion[2]
obs["imu.quat.z"] = imu.quaternion[3]
if imu.rpy:
obs["imu.rpy.roll"] = imu.rpy[0]
obs["imu.rpy.pitch"] = imu.rpy[1]
obs["imu.rpy.yaw"] = imu.rpy[2]
wr = getattr(lowstate, "wireless_remote", None)
if wr:
obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
return obs
def make_locomotion_controller(name: str | None):
"""Instantiate a locomotion controller by class name. Returns None if name is None."""
if name is None:
@@ -112,6 +112,24 @@ class SmplStream:
"""True once at least one real frame has been received."""
return self._got_first
@property
def seconds_since_last(self) -> float:
"""Wall-clock seconds since the last real frame (inf before the first)."""
if not self._got_first:
return float("inf")
return time.time() - self._last_recv_t
@property
def is_stale(self) -> bool:
"""True when the stream has gone silent past ``stale_after_s``.
Consumers use this to stop feeding a frozen pose and let the controller
fall back to a safe standing/locomotion mode.
"""
if not self._got_first or not self.stale_after_s:
return False
return self.seconds_since_last > self.stale_after_s
def reset(self):
self._buf.clear()
self._got_first = False
+3 -38
View File
@@ -38,6 +38,7 @@ from .g1_utils import (
G1_29_JointArmIndex,
G1_29_JointIndex,
default_remote_input,
lowstate_to_obs,
make_locomotion_controller,
)
@@ -428,44 +429,8 @@ class UnitreeG1(Robot):
if lowstate is None:
return {}
obs = {}
# Motors - q, dq, tau for all joints
for motor in G1_29_JointIndex:
name = motor.name
idx = motor.value
obs[f"{name}.q"] = lowstate.motor_state[idx].q
obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
# IMU - gyroscope
if lowstate.imu_state.gyroscope:
obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
# IMU - accelerometer
if lowstate.imu_state.accelerometer:
obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
# IMU - quaternion
if lowstate.imu_state.quaternion:
obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
# IMU - rpy
if lowstate.imu_state.rpy:
obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
# Wireless remote (raw bytes for teleoperator)
if lowstate.wireless_remote:
obs["wireless_remote"] = lowstate.wireless_remote
# Motors + IMU + wireless remote (shared lowstate -> obs mapping)
obs = lowstate_to_obs(lowstate)
# Cameras - read images from ZMQ cameras
for cam_name, cam in self._cameras.items():
@@ -91,9 +91,11 @@ class PicoHeadset(Teleoperator):
if self._stream is None:
raise RuntimeError(f"{self} is not connected")
window = self._stream.step()
# Hold back until the headset is actually streaming, so the controller does
# not switch to whole-body tracking of an all-zero (collapsed) pose.
if not self._stream.has_data:
# Emit SMPL only while the headset is actively streaming: hold back before
# the first frame (so the controller doesn't track an all-zero collapsed
# pose) and once the stream goes stale (so the controller falls back to a
# safe standing/locomotion mode instead of freezing on the last pose).
if not self._stream.has_data or self._stream.is_stale:
return {}
return {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)}
@@ -0,0 +1,153 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Standalone ``rt/smpl`` publisher for the PICO headset (no gear_sonic / torch).
Reads 24 body-joint poses from the XRoboToolkit SDK, runs pure-numpy SMPL forward
kinematics + canonicalization (``smpl_fk.py``), and publishes one canonical
``(24, 3)`` SMPL frame per tick over ZMQ on the ``rt/smpl`` topic the exact
message ``lerobot.robots.unitree_g1.smpl_stream.SmplStream`` consumes.
This makes the LeRobot side self-contained: the only runtime dependency to drive
SONIC whole-body teleop from the headset is the ``xrobotoolkit_sdk`` Python package
(plus numpy/scipy/pyzmq), not the full ``gear_sonic`` stack.
Usage:
# Real headset (XRoboToolkit PC Service must be running and connected):
python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --port 5560
# No hardware — emit a synthetic waving motion to test the consumer end-to-end:
python -m lerobot.teleoperators.pico_headset.pico_publisher --fake
"""
from __future__ import annotations
import argparse
import contextlib
import json
import time
import numpy as np
import zmq
from lerobot.teleoperators.pico_headset.smpl_fk import SmplForwardKinematics
SMPL_TOPIC = "rt/smpl"
DEFAULT_SMPL_PORT = 5560
def pack_message(
smpl_joints_local: np.ndarray,
frame_index: int,
stamp_ns: int,
root_quat: np.ndarray | None = None,
root_transl: np.ndarray | None = None,
) -> bytes:
"""Build the rt/smpl JSON message (single frame, topic embedded in payload)."""
data = {
"smpl_joints_local": np.asarray(smpl_joints_local, np.float32).reshape(-1).tolist(),
"frame_index": int(frame_index),
"stamp_ns": int(stamp_ns),
}
if root_quat is not None:
data["root_quat"] = np.asarray(root_quat, np.float32).reshape(-1).tolist()
if root_transl is not None:
data["root_transl"] = np.asarray(root_transl, np.float32).reshape(-1).tolist()
return json.dumps({"topic": SMPL_TOPIC, "data": data}).encode("utf-8")
def _fake_body_poses(t: float) -> np.ndarray:
"""Synthetic (24, 7) body poses: identity rotations + a gently waving right arm."""
poses = np.zeros((24, 7), np.float32)
poses[:, 6] = 1.0 # unit quaternion (qw = 1), scalar-last
poses[:, 1] = 1.0 # ~1 m pelvis height (positions only matter for root_transl)
# Wave the right shoulder (SMPL body joint 17) about Z.
ang = 0.5 * np.sin(2.0 * np.pi * 0.5 * t)
poses[17, 3:7] = [0.0, 0.0, np.sin(ang / 2), np.cos(ang / 2)]
return poses
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--port", type=int, default=DEFAULT_SMPL_PORT, help="ZMQ PUB port for rt/smpl")
p.add_argument("--fps", type=float, default=50.0, help="Target publish rate (Hz)")
p.add_argument("--skeleton", type=str, default=None, help="Path to smpl_skeleton.npz")
p.add_argument("--fake", action="store_true", help="Publish synthetic motion (no headset)")
args = p.parse_args()
fk = SmplForwardKinematics(args.skeleton) if args.skeleton else SmplForwardKinematics()
xrt = None
if not args.fake:
try:
import xrobotoolkit_sdk as xrt # noqa: PLC0415
except ImportError as e:
raise SystemExit(
"xrobotoolkit_sdk not available. Install it, or run with --fake to test "
"the pipeline without a headset."
) from e
xrt.init()
print("[pico_publisher] XRoboToolkit initialized")
ctx = zmq.Context.instance()
sock = ctx.socket(zmq.PUB)
sock.bind(f"tcp://*:{args.port}")
print(f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz")
period = 1.0 / max(1.0, args.fps)
frame_index = 0
t0 = time.time()
try:
while True:
loop_start = time.time()
if args.fake:
body_poses = _fake_body_poses(loop_start - t0)
stamp_ns = time.time_ns()
else:
body_poses = np.asarray(xrt.get_body_joints_pose(), np.float32)
stamp_ns = int(xrt.get_time_stamp_ns())
if body_poses.shape != (24, 7):
time.sleep(0.005)
continue
out = fk.compute(body_poses)
sock.send(
pack_message(
out["smpl_joints_local"],
frame_index,
stamp_ns,
root_quat=out["root_quat"],
root_transl=out["root_transl"],
)
)
frame_index += 1
if frame_index % int(max(1, args.fps)) == 0:
print(f"[pico_publisher] sent {frame_index} frames", end="\r")
dt = time.time() - loop_start
if dt < period:
time.sleep(period - dt)
except KeyboardInterrupt:
print("\n[pico_publisher] stopping")
finally:
sock.close(0)
if xrt is not None:
with contextlib.suppress(Exception):
xrt.close()
if __name__ == "__main__":
main()
@@ -0,0 +1,194 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Standalone SMPL forward kinematics + canonicalization in pure numpy/scipy.
This mirrors the ``rt/smpl`` producer path in ``gear_sonic`` (``compute_from_body_poses``
-> ``process_smpl_joints`` -> ``compute_human_joints``) without depending on torch,
the SMPL mesh model, or ``gear_sonic``. It only needs a small fixed skeleton table
(rest-pose joints + kinematic tree), vendored in ``assets/smpl_skeleton.npz``.
Given the 24 body-joint poses reported by the XRoboToolkit headset SDK
(``xrt.get_body_joints_pose()`` -> (24, 7) of ``[x, y, z, qx, qy, qz, qw]``), it
produces the root-orientation-removed 24x3 SMPL joints the SONIC encoder expects,
plus the root orientation quaternion and pelvis translation.
Quaternions are scalar-first (w, x, y, z) unless noted.
"""
from pathlib import Path
import numpy as np
from scipy.spatial.transform import Rotation as R
# 24-joint parent tree used by the headset body-pose stream (SMPL-X body subset).
# Matches PoseStreamer.parent_indices in gear_sonic's pico_manager_thread_server.py.
BODY24_PARENTS = np.array(
[-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19, 20, 22],
dtype=np.int64,
)
# FK output joints: first 22 SMPL body joints + two thumb tips (SMPL-X indices 39, 54).
OUTPUT_JOINT_INDEX = np.concatenate([np.arange(22), np.array([39, 54])])
_SKELETON_PATH = Path(__file__).parent / "assets" / "smpl_skeleton.npz"
# ── quaternion helpers (scalar-first w, x, y, z) ─────────────────────────────
def aa_to_quat(aa: np.ndarray) -> np.ndarray:
aa = np.asarray(aa, np.float64)
angle = np.linalg.norm(aa, axis=-1, keepdims=True)
small = angle < 1e-8
safe = np.where(small, 1.0, angle)
axis = np.where(small, 0.0, aa / safe)
half = angle * 0.5
return np.concatenate([np.cos(half), axis * np.sin(half)], axis=-1)
def quat_to_aa(q: np.ndarray) -> np.ndarray:
q = np.asarray(q, np.float64)
w = q[..., 0:1]
xyz = q[..., 1:]
n = np.linalg.norm(xyz, axis=-1, keepdims=True)
angle = 2.0 * np.arctan2(n, w)
small = n < 1e-8
safe = np.where(small, 1.0, n)
axis = np.where(small, 0.0, xyz / safe)
return axis * angle
def quat_mul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3]
bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3]
return np.stack(
[
aw * bw - ax * bx - ay * by - az * bz,
aw * bx + ax * bw + ay * bz - az * by,
aw * by - ax * bz + ay * bw + az * bx,
aw * bz + ax * by - ay * bx + az * bw,
],
axis=-1,
)
def quat_conj(q: np.ndarray) -> np.ndarray:
return np.concatenate([q[..., 0:1], -q[..., 1:]], axis=-1)
def quat_inv(q: np.ndarray) -> np.ndarray:
c = quat_conj(q)
return c / (np.linalg.norm(c, axis=-1, keepdims=True) + 1e-12)
def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray:
xyz = q[..., 1:]
w = q[..., 0:1]
t = 2.0 * np.cross(xyz, v)
return v + w * t + np.cross(xyz, t)
def smpl_root_ytoz_up(root_quat: np.ndarray) -> np.ndarray:
"""Rotate the root quaternion 90 deg about X to map SMPL Y-up to robot Z-up."""
base = aa_to_quat(np.array([np.pi / 2.0, 0.0, 0.0]))
return quat_mul(base, root_quat)
def remove_smpl_base_rot(root_quat: np.ndarray) -> np.ndarray:
"""Conjugate out SMPL's default rest orientation ([0.5, 0.5, 0.5, 0.5])."""
base_conj = quat_conj(np.array([0.5, 0.5, 0.5, 0.5]))
return quat_mul(root_quat, base_conj)
# ── forward kinematics ───────────────────────────────────────────────────────
class SmplForwardKinematics:
"""Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""
def __init__(self, skeleton_path: str | Path = _SKELETON_PATH):
data = np.load(skeleton_path)
self.J = data["J"].astype(np.float64) # (55, 3) rest joint positions
self.parents = data["parents"].astype(np.int64) # (55,) kinematic tree
self.n_joints = self.J.shape[0]
def _fk(self, full_pose_aa: np.ndarray) -> np.ndarray:
"""full_pose_aa: (n_joints, 3) axis-angle (joint 0 = global). Returns (24, 3)."""
rot = R.from_rotvec(full_pose_aa).as_matrix() # (n, 3, 3)
rel = self.J.copy()
rel[1:] -= self.J[self.parents[1:]]
transforms = np.zeros((self.n_joints, 4, 4), np.float64)
transforms[:, :3, :3] = rot
transforms[:, :3, 3] = rel
transforms[:, 3, 3] = 1.0
chain = [transforms[0]]
for i in range(1, self.n_joints):
chain.append(chain[self.parents[i]] @ transforms[i])
joints = np.stack(chain)[:, :3, 3]
return joints[OUTPUT_JOINT_INDEX]
def compute(self, body_poses_np: np.ndarray) -> dict:
"""Convert (24, 7) headset body poses to canonical SMPL joints.
Args:
body_poses_np: (24, 7) rows of [x, y, z, qx, qy, qz, qw] (scalar-last).
Returns:
dict with:
- smpl_joints_local: (24, 3) root-orientation-removed joints
- root_quat: (4,) root/torso orientation (w, x, y, z)
- root_transl: (3,) pelvis translation
"""
body_poses_np = np.asarray(body_poses_np, np.float64)
positions = body_poses_np[:, :3]
# Global joint rotations from the headset (scalar-last), with the SMPL
# +180 deg-about-Y frame fix, converted to per-joint local axis-angle.
global_rots = R.from_quat(body_poses_np[:, 3:7]) * R.from_euler("y", 180, degrees=True)
gm = global_rots.as_matrix() # (24, 3, 3)
local_aa = np.zeros((24, 3), np.float64)
for i in range(24):
p = BODY24_PARENTS[i]
m = gm[i] if p == -1 else gm[p].T @ gm[i]
local_aa[i] = R.from_matrix(m).as_rotvec()
global_orient = local_aa[0]
body_pose = local_aa[1:].reshape(-1)[:63] # 21 body joints
# Root: Y-up -> Z-up, then run FK with the transformed root.
root_quat = aa_to_quat(global_orient)
root_quat = smpl_root_ytoz_up(root_quat)
global_orient_new = quat_to_aa(root_quat)
full_pose = np.concatenate(
[global_orient_new, body_pose, np.zeros(3 * self.n_joints - 66)]
).reshape(self.n_joints, 3)
joints = self._fk(full_pose) # (24, 3)
# Canonicalize: strip SMPL base rot and the root orientation.
root_quat = remove_smpl_base_rot(root_quat)
inv = np.broadcast_to(quat_inv(root_quat), (joints.shape[0], 4))
smpl_joints_local = quat_apply(inv, joints)
return {
"smpl_joints_local": smpl_joints_local.astype(np.float32),
"root_quat": root_quat.astype(np.float32),
"root_transl": positions[0].astype(np.float32),
}