g1-teleop/scripts/body_tracker.gd

extends Node
## Reads XR_FB_body_tracking joints each frame via Godot's XRBodyTracker.
## Computes chest-relative wrist positions and emits tracking data.
##
## Meta body tracking provides 70 joints. We use:
##   CHEST (5)  - torso orientation (solves body rotation problem)
##   HEAD (7)   - head pose
##   LEFT_HAND_WRIST (12)  - left wrist
##   RIGHT_HAND_WRIST (17) - right wrist
##   Hand joints (18-69)   - finger positions
##
## Output poses are in Godot's coordinate system (Y-up, -Z forward).
## The server handles conversion to robot conventions.

## Emitted every frame with tracking data dict ready to send.
signal tracking_data_ready(data: Dictionary)

## Joint indices from XR_FB_body_tracking
## Reference: Meta OpenXR body tracking extension
const JOINT_HIPS := 1
const JOINT_SPINE_LOWER := 2
const JOINT_SPINE_MIDDLE := 3
const JOINT_SPINE_UPPER := 4
const JOINT_CHEST := 5
const JOINT_NECK := 6
const JOINT_HEAD := 7

const JOINT_LEFT_SHOULDER := 8
const JOINT_LEFT_SCAPULA := 9
const JOINT_LEFT_ARM_UPPER := 10
const JOINT_LEFT_ARM_LOWER := 11
const JOINT_LEFT_HAND_WRIST := 12

const JOINT_RIGHT_SHOULDER := 13
const JOINT_RIGHT_SCAPULA := 14
const JOINT_RIGHT_ARM_UPPER := 15
const JOINT_RIGHT_ARM_LOWER := 16
const JOINT_RIGHT_HAND_WRIST := 17

# Hand joints start at index 18 for left hand, 43 for right hand
# 25 joints per hand (same layout as XR_EXT_hand_tracking)
const JOINT_LEFT_HAND_START := 18
const JOINT_RIGHT_HAND_START := 43
const HAND_JOINT_COUNT := 25

## Total body joint count
const BODY_JOINT_COUNT := 70

## Tracking state
var body_tracker_name: StringName = &"/user/body_tracker"
var is_tracking: bool = false
var frames_since_last_send: int = 0

## Send rate control: target ~30 Hz (every other frame at 72 Hz)
@export var send_every_n_frames: int = 2

## Debug logging
@export var debug_log: bool = false
var _log_counter: int = 0


func _ready() -> void:
	print("[BodyTracker] Initialized, waiting for body tracking data...")


func _process(_delta: float) -> void:
	frames_since_last_send += 1
	if frames_since_last_send < send_every_n_frames:
		return
	frames_since_last_send = 0

	var tracker := XRServer.get_tracker(body_tracker_name) as XRBodyTracker
	if tracker == null:
		if is_tracking:
			print("[BodyTracker] Lost body tracking")
			is_tracking = false
		return

	if not tracker.get_has_tracking_data():
		if is_tracking:
			print("[BodyTracker] Body tracking data unavailable")
			is_tracking = false
		return

	if not is_tracking:
		print("[BodyTracker] Body tracking active!")
		is_tracking = true

	# Read key joint poses
	var chest_xform := tracker.get_joint_transform(JOINT_CHEST)
	var head_xform := tracker.get_joint_transform(JOINT_HEAD)
	var left_wrist_xform := tracker.get_joint_transform(JOINT_LEFT_HAND_WRIST)
	var right_wrist_xform := tracker.get_joint_transform(JOINT_RIGHT_HAND_WRIST)

	# Compute chest-relative wrist positions
	# This is the key advantage: eliminates head-position subtraction artifacts
	var chest_inv := chest_xform.affine_inverse()
	var left_wrist_rel := chest_inv * left_wrist_xform
	var right_wrist_rel := chest_inv * right_wrist_xform

	# Build tracking data packet
	var data := {}
	data["type"] = "tracking"
	data["timestamp"] = Time.get_ticks_msec()

	# Head pose (world space) - 7 floats: pos(3) + quat(4)
	data["head"] = _xform_to_pose7(head_xform)

	# Chest pose (world space) - 7 floats: pos(3) + quat(4)
	data["chest"] = _xform_to_pose7(chest_xform)

	# Chest-relative wrist poses as 4x4 matrices (16 floats each, column-major)
	# These replace the world-space wrist poses; no head subtraction needed
	data["left_wrist"] = _xform_to_mat16(left_wrist_rel)
	data["right_wrist"] = _xform_to_mat16(right_wrist_rel)

	# Hand joint positions (25 joints per hand, relative to wrist)
	var left_hand_pos := _get_hand_positions(tracker, JOINT_LEFT_HAND_START, left_wrist_xform)
	var right_hand_pos := _get_hand_positions(tracker, JOINT_RIGHT_HAND_START, right_wrist_xform)
	data["left_hand_pos"] = left_hand_pos
	data["right_hand_pos"] = right_hand_pos

	# Hand joint rotations (25 joints per hand, relative to wrist)
	var left_hand_rot := _get_hand_rotations(tracker, JOINT_LEFT_HAND_START, left_wrist_xform)
	var right_hand_rot := _get_hand_rotations(tracker, JOINT_RIGHT_HAND_START, right_wrist_xform)
	data["left_hand_rot"] = left_hand_rot
	data["right_hand_rot"] = right_hand_rot

	# Debug logging every ~2 seconds
	if debug_log:
		_log_counter += 1
		if _log_counter >= 60:
			_log_counter = 0
			print("[BodyTracker] chest=", chest_xform.origin,
				  " head=", head_xform.origin,
				  " L_wrist_rel=", left_wrist_rel.origin,
				  " R_wrist_rel=", right_wrist_rel.origin)

	tracking_data_ready.emit(data)


## Convert Transform3D to 7-float pose: [x, y, z, qx, qy, qz, qw]
func _xform_to_pose7(xform: Transform3D) -> Array:
	var pos := xform.origin
	var quat := xform.basis.get_rotation_quaternion()
	return [pos.x, pos.y, pos.z, quat.x, quat.y, quat.z, quat.w]


## Convert Transform3D to 16-float column-major 4x4 matrix
## Layout: [m00, m10, m20, 0, m01, m11, m21, 0, m02, m12, m22, 0, tx, ty, tz, 1]
## This matches NumPy's Fortran ('F') order used by the existing teleop code.
func _xform_to_mat16(xform: Transform3D) -> Array:
	var b := xform.basis
	var o := xform.origin
	# Godot Basis: b[col][row] — but we store column-major
	# Column 0: basis x-axis
	# Column 1: basis y-axis
	# Column 2: basis z-axis
	# Column 3: origin (translation)
	return [
		b.x.x, b.x.y, b.x.z, 0.0,  # Column 0
		b.y.x, b.y.y, b.y.z, 0.0,  # Column 1
		b.z.x, b.z.y, b.z.z, 0.0,  # Column 2
		o.x,   o.y,   o.z,   1.0   # Column 3
	]


## Get 25 hand joint positions relative to wrist, as flat array (75 floats)
## Each joint: [x, y, z] relative to wrist
func _get_hand_positions(tracker: XRBodyTracker, start_idx: int, wrist_xform: Transform3D) -> Array:
	var wrist_inv := wrist_xform.affine_inverse()
	var positions := []
	positions.resize(HAND_JOINT_COUNT * 3)
	for i in range(HAND_JOINT_COUNT):
		var joint_xform := tracker.get_joint_transform(start_idx + i)
		var rel := wrist_inv * joint_xform
		positions[i * 3 + 0] = rel.origin.x
		positions[i * 3 + 1] = rel.origin.y
		positions[i * 3 + 2] = rel.origin.z
	return positions


## Get 25 hand joint rotations relative to wrist, as flat array (225 floats)
## Each joint: 9 floats (3x3 rotation matrix, column-major)
func _get_hand_rotations(tracker: XRBodyTracker, start_idx: int, wrist_xform: Transform3D) -> Array:
	var wrist_rot_inv := wrist_xform.basis.inverse()
	var rotations := []
	rotations.resize(HAND_JOINT_COUNT * 9)
	for i in range(HAND_JOINT_COUNT):
		var joint_xform := tracker.get_joint_transform(start_idx + i)
		var rel_basis := wrist_rot_inv * joint_xform.basis
		# Store as column-major 3x3
		rotations[i * 9 + 0] = rel_basis.x.x
		rotations[i * 9 + 1] = rel_basis.x.y
		rotations[i * 9 + 2] = rel_basis.x.z
		rotations[i * 9 + 3] = rel_basis.y.x
		rotations[i * 9 + 4] = rel_basis.y.y
		rotations[i * 9 + 5] = rel_basis.y.z
		rotations[i * 9 + 6] = rel_basis.z.x
		rotations[i * 9 + 7] = rel_basis.z.y
		rotations[i * 9 + 8] = rel_basis.z.z
	return rotations