microduck-electric-slide-policy / code /physics_choreography.py

Add files using upload-large-folder tool

32502a8 verified about 2 months ago

6.05 kB

"""Electric Slide commands for the official physics-based walking policy."""

from __future__ import annotations

from dataclasses import dataclass
import math
from pathlib import Path

import numpy as np

from choreography import BEAT_SECONDS, COUNTS_PER_WALL, SHOW_BEATS, SHOW_SECONDS


@dataclass(frozen=True)
class PolicyCommand:
    twist: tuple[float, float, float]
    head_pose: tuple[float, float, float, float]
    body_pose: tuple[float, float, float, float, float, float]


# Body-frame velocity commands for each count: forward, lateral, yaw rate.
# Pollen's current simulator deliberately leaves the lateral slot at zero,
# and this published policy does not produce a backward gait.  The dance
# therefore uses supported forward footfalls plus mirrored arcs for its
# right/left phrases.  The second half bends through one physical wall turn.
COUNT_TWISTS = (
    (0.22, 0.00, 0.65),
    (0.22, 0.00, 0.65),
    (0.22, 0.00, 0.65),
    (0.00, 0.00, 0.00),
    (0.22, 0.00, -0.65),
    (0.22, 0.00, -0.65),
    (0.22, 0.00, -0.65),
    (0.00, 0.00, 0.00),
    (0.18, 0.00, 1.00),
    (0.18, 0.00, 1.00),
    (0.18, 0.00, 1.00),
    (0.00, 0.00, 1.00),
    (0.00, 0.00, 0.00),
    (0.14, 0.00, 1.00),
    (0.00, 0.00, 0.00),
    (0.10, 0.00, 1.00),
    (0.10, 0.00, 1.00),
    (0.10, 0.00, 1.00),
)

MOTION_TIMELINE = (
    Path(__file__).resolve().parent
    / "generated/learned_motion_retarget_50hz.npz"
)
_TIMELINE_CACHE: dict[str, np.ndarray] | None = None


def _timeline() -> dict[str, np.ndarray] | None:
    global _TIMELINE_CACHE
    if _TIMELINE_CACHE is None and MOTION_TIMELINE.exists():
        data = np.load(MOTION_TIMELINE, allow_pickle=True)
        _TIMELINE_CACHE = {name: data[name] for name in data.files}
    return _TIMELINE_CACHE


def _motion_sample(time_seconds: float) -> tuple[np.ndarray, int] | None:
    data = _timeline()
    if data is None:
        return None
    commands = data["command_13d"]
    control_hz = float(data["control_hz"])
    phase = (time_seconds % SHOW_SECONDS) * control_hz
    lower = int(math.floor(phase)) % len(commands)
    upper = (lower + 1) % len(commands)
    blend = phase - math.floor(phase)
    command = commands[lower] * (1.0 - blend) + commands[upper] * blend
    policy_code = int(data["policy_code"][lower])
    return command, policy_code


def _smoothstep(value: float) -> float:
    value = max(0.0, min(1.0, value))
    return value * value * (3.0 - 2.0 * value)


def _blended_twist(beat: float) -> tuple[float, float, float]:
    count = int(math.floor(beat)) % COUNTS_PER_WALL
    phase = beat - math.floor(beat)
    current = COUNT_TWISTS[count]
    previous = COUNT_TWISTS[(count - 1) % COUNTS_PER_WALL]
    # Do not carry the prior wall's turn into the next grapevine.
    if count == 0:
        previous = current
    blend = _smoothstep(phase / 0.10)
    return tuple(previous[i] + (current[i] - previous[i]) * blend for i in range(3))


def command_for_time(time_seconds: float, duck_index: int = 0) -> PolicyCommand:
    """Return the 61D-policy command blocks for one show instant."""
    motion = _motion_sample(time_seconds)
    if motion is not None:
        command, _ = motion
        return PolicyCommand(
            twist=tuple(float(value) for value in command[0:3]),
            head_pose=tuple(float(value) for value in command[3:7]),
            body_pose=tuple(float(value) for value in command[7:13]),
        )
    beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
    count = int(math.floor(beat)) % COUNTS_PER_WALL
    phase = beat - math.floor(beat)
    pulse = math.sin(math.pi * phase)
    # Keep the physics commands exactly synchronized.  Even small pose-command
    # phase offsets can amplify into formation drift through the learned gait.
    flock_phase = 0.0
    groove = math.sin(2.0 * math.pi * (beat / 2.0 + flock_phase))

    twist = _blended_twist(beat)

    # Head attitude stays within the command ranges used in training.
    neck_pitch = 0.05 * groove
    head_pitch = -0.05 * groove
    head_yaw = 0.10 * math.sin(2.0 * math.pi * (beat / 4.0 + flock_phase))
    head_roll = 0.09 * groove

    body_z = -0.004 * abs(groove)
    body_roll = 0.035 * groove
    body_pitch = 0.018 * math.sin(2.0 * math.pi * beat / 4.0)

    # Character section: grounded scoop, relaxed lean, then head leads turn.
    if count in (12, 13):
        amount = _smoothstep(phase) if count == 12 else pulse
        body_z -= 0.010 * amount
        body_pitch += 0.035 * amount
        neck_pitch += 0.08 * amount
        head_pitch -= 0.06 * amount
        head_yaw += 0.08 * amount
    elif count in (14, 15):
        amount = _smoothstep(phase) if count == 14 else pulse
        body_pitch -= 0.035 * amount
        neck_pitch -= 0.06 * amount
        head_pitch += 0.05 * amount
    elif count in (16, 17):
        head_yaw += 0.20 * pulse
        head_roll += 0.04 * groove

    return PolicyCommand(
        twist=twist,
        head_pose=(neck_pitch, head_pitch, head_yaw, head_roll),
        body_pose=(0.0, 0.0, body_z, body_roll, body_pitch, 0.0),
    )


def policy_for_time(time_seconds: float) -> str:
    """Choose official one-shot kick policies for the two accent counts."""
    motion = _motion_sample(time_seconds)
    if motion is not None:
        _, policy_code = motion
        return ("walk", "kick_left", "kick_right")[policy_code]
    beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
    count = int(math.floor(beat)) % COUNTS_PER_WALL
    if count == 12:
        return "kick_left"
    if count == 14:
        return "kick_right"
    return "walk"


def integrated_turn_radians(samples: int = 10000) -> float:
    """Numerically integrate one phrase's yaw command for a regression test."""
    duration = COUNTS_PER_WALL * BEAT_SECONDS
    dt = duration / samples
    return sum(command_for_time((index + 0.5) * dt).twist[2] * dt for index in range(samples))


__all__ = [
    "COUNT_TWISTS",
    "PolicyCommand",
    "SHOW_SECONDS",
    "command_for_time",
    "integrated_turn_radians",
    "policy_for_time",
]

	"""Electric Slide commands for the official physics-based walking policy."""

	from __future__ import annotations

	from dataclasses import dataclass
	import math
	from pathlib import Path

	import numpy as np

	from choreography import BEAT_SECONDS, COUNTS_PER_WALL, SHOW_BEATS, SHOW_SECONDS


	@dataclass(frozen=True)
	class PolicyCommand:
	twist: tuple[float, float, float]
	head_pose: tuple[float, float, float, float]
	body_pose: tuple[float, float, float, float, float, float]


	# Body-frame velocity commands for each count: forward, lateral, yaw rate.
	# Pollen's current simulator deliberately leaves the lateral slot at zero,
	# and this published policy does not produce a backward gait. The dance
	# therefore uses supported forward footfalls plus mirrored arcs for its
	# right/left phrases. The second half bends through one physical wall turn.
	COUNT_TWISTS = (
	(0.22, 0.00, 0.65),
	(0.22, 0.00, 0.65),
	(0.22, 0.00, 0.65),
	(0.00, 0.00, 0.00),
	(0.22, 0.00, -0.65),
	(0.22, 0.00, -0.65),
	(0.22, 0.00, -0.65),
	(0.00, 0.00, 0.00),
	(0.18, 0.00, 1.00),
	(0.18, 0.00, 1.00),
	(0.18, 0.00, 1.00),
	(0.00, 0.00, 1.00),
	(0.00, 0.00, 0.00),
	(0.14, 0.00, 1.00),
	(0.00, 0.00, 0.00),
	(0.10, 0.00, 1.00),
	(0.10, 0.00, 1.00),
	(0.10, 0.00, 1.00),
	)

	MOTION_TIMELINE = (
	Path(__file__).resolve().parent
	/ "generated/learned_motion_retarget_50hz.npz"
	)
	_TIMELINE_CACHE: dict[str, np.ndarray] \| None = None


	def _timeline() -> dict[str, np.ndarray] \| None:
	global _TIMELINE_CACHE
	if _TIMELINE_CACHE is None and MOTION_TIMELINE.exists():
	data = np.load(MOTION_TIMELINE, allow_pickle=True)
	_TIMELINE_CACHE = {name: data[name] for name in data.files}
	return _TIMELINE_CACHE


	def _motion_sample(time_seconds: float) -> tuple[np.ndarray, int] \| None:
	data = _timeline()
	if data is None:
	return None
	commands = data["command_13d"]
	control_hz = float(data["control_hz"])
	phase = (time_seconds % SHOW_SECONDS) * control_hz
	lower = int(math.floor(phase)) % len(commands)
	upper = (lower + 1) % len(commands)
	blend = phase - math.floor(phase)
	command = commands[lower] * (1.0 - blend) + commands[upper] * blend
	policy_code = int(data["policy_code"][lower])
	return command, policy_code


	def _smoothstep(value: float) -> float:
	value = max(0.0, min(1.0, value))
	return value * value * (3.0 - 2.0 * value)


	def _blended_twist(beat: float) -> tuple[float, float, float]:
	count = int(math.floor(beat)) % COUNTS_PER_WALL
	phase = beat - math.floor(beat)
	current = COUNT_TWISTS[count]
	previous = COUNT_TWISTS[(count - 1) % COUNTS_PER_WALL]
	# Do not carry the prior wall's turn into the next grapevine.
	if count == 0:
	previous = current
	blend = _smoothstep(phase / 0.10)
	return tuple(previous[i] + (current[i] - previous[i]) * blend for i in range(3))


	def command_for_time(time_seconds: float, duck_index: int = 0) -> PolicyCommand:
	"""Return the 61D-policy command blocks for one show instant."""
	motion = _motion_sample(time_seconds)
	if motion is not None:
	command, _ = motion
	return PolicyCommand(
	twist=tuple(float(value) for value in command[0:3]),
	head_pose=tuple(float(value) for value in command[3:7]),
	body_pose=tuple(float(value) for value in command[7:13]),
	)
	beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
	count = int(math.floor(beat)) % COUNTS_PER_WALL
	phase = beat - math.floor(beat)
	pulse = math.sin(math.pi * phase)
	# Keep the physics commands exactly synchronized. Even small pose-command
	# phase offsets can amplify into formation drift through the learned gait.
	flock_phase = 0.0
	groove = math.sin(2.0 * math.pi * (beat / 2.0 + flock_phase))

	twist = _blended_twist(beat)

	# Head attitude stays within the command ranges used in training.
	neck_pitch = 0.05 * groove
	head_pitch = -0.05 * groove
	head_yaw = 0.10 * math.sin(2.0 * math.pi * (beat / 4.0 + flock_phase))
	head_roll = 0.09 * groove

	body_z = -0.004 * abs(groove)
	body_roll = 0.035 * groove
	body_pitch = 0.018 * math.sin(2.0 * math.pi * beat / 4.0)

	# Character section: grounded scoop, relaxed lean, then head leads turn.
	if count in (12, 13):
	amount = _smoothstep(phase) if count == 12 else pulse
	body_z -= 0.010 * amount
	body_pitch += 0.035 * amount
	neck_pitch += 0.08 * amount
	head_pitch -= 0.06 * amount
	head_yaw += 0.08 * amount
	elif count in (14, 15):
	amount = _smoothstep(phase) if count == 14 else pulse
	body_pitch -= 0.035 * amount
	neck_pitch -= 0.06 * amount
	head_pitch += 0.05 * amount
	elif count in (16, 17):
	head_yaw += 0.20 * pulse
	head_roll += 0.04 * groove

	return PolicyCommand(
	twist=twist,
	head_pose=(neck_pitch, head_pitch, head_yaw, head_roll),
	body_pose=(0.0, 0.0, body_z, body_roll, body_pitch, 0.0),
	)


	def policy_for_time(time_seconds: float) -> str:
	"""Choose official one-shot kick policies for the two accent counts."""
	motion = _motion_sample(time_seconds)
	if motion is not None:
	_, policy_code = motion
	return ("walk", "kick_left", "kick_right")[policy_code]
	beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
	count = int(math.floor(beat)) % COUNTS_PER_WALL
	if count == 12:
	return "kick_left"
	if count == 14:
	return "kick_right"
	return "walk"


	def integrated_turn_radians(samples: int = 10000) -> float:
	"""Numerically integrate one phrase's yaw command for a regression test."""
	duration = COUNTS_PER_WALL * BEAT_SECONDS
	dt = duration / samples
	return sum(command_for_time((index + 0.5) * dt).twist[2] * dt for index in range(samples))


	__all__ = [
	"COUNT_TWISTS",
	"PolicyCommand",
	"SHOW_SECONDS",
	"command_for_time",
	"integrated_turn_radians",
	"policy_for_time",
	]