Add files using upload-large-folder tool

32502a8 verified about 2 months ago

6.8 kB

"""Music-timed, 50 Hz Electric Slide choreography for a Microduck flock.
The motion is deliberately kinematic for the first visual prototype: root poses
and servo poses are sampled directly in MuJoCo. That makes the 25-duck show
repeatable while keeping a clean seam for replacing the sampler with an ONNX
policy later.
"""

from __future__ import annotations

from dataclasses import dataclass
import math


BPM = 111.0
BEAT_SECONDS = 60.0 / BPM
COUNTS_PER_WALL = 18
WALLS = 4
SHOW_BEATS = COUNTS_PER_WALL * WALLS
SHOW_SECONDS = SHOW_BEATS * BEAT_SECONDS
CONTROL_HZ = 50

# Official STAND keyframe, in official actuator/joint order.
STAND_POSE = (
    0.0,
    -0.08726646259971647,
    -0.457924,
    -0.004940,
    0.452984,
    0.3490658503988659,
    0.3490658503988659,
    0.0,
    0.0,
    0.0,
    0.08726646259971647,
    0.457924,
    0.004940,
    -0.452984,
)

# Per-count root translation in the duck's local frame (forward, left) and yaw.
# Right is negative local y. Four 18-count walls trace a closed square.
ROOT_STEPS = (
    (0.00, -0.10, 0.0),  # grapevine right: side
    (0.00, -0.04, 0.0),  # cross behind
    (0.00, -0.10, 0.0),  # side
    (0.00, 0.00, 0.0),   # touch
    (0.00, 0.10, 0.0),   # grapevine left: side
    (0.00, 0.04, 0.0),   # cross behind
    (0.00, 0.10, 0.0),   # side
    (0.00, 0.00, 0.0),   # touch
    (-0.08, 0.00, 0.0),  # walk back
    (-0.08, 0.00, 0.0),
    (-0.08, 0.00, 0.0),
    (0.00, 0.00, 0.0),   # touch
    (0.08, 0.00, 0.0),   # rock forward
    (0.00, 0.00, 0.0),   # toe touch + scoop
    (-0.08, 0.00, 0.0),  # rock back
    (0.00, 0.00, 0.0),   # toe touch + lean
    (0.05, 0.00, 0.0),   # step into turn
    (0.00, 0.00, math.pi / 2.0),  # scuff + quarter turn left
)


@dataclass(frozen=True)
class DuckPose:
    position: tuple[float, float, float]
    quaternion: tuple[float, float, float, float]
    joints: tuple[float, ...]


def formation_positions(grid: int = 5, spacing: float = 0.52) -> list[tuple[float, float]]:
    """Return a centered, row-major square formation."""
    center = (grid - 1) / 2.0
    return [
        ((col - center) * spacing, (center - row) * spacing)
        for row in range(grid)
        for col in range(grid)
    ]


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


def _rotate(x: float, y: float, yaw: float) -> tuple[float, float]:
    cosine, sine = math.cos(yaw), math.sin(yaw)
    return cosine * x - sine * y, sine * x + cosine * y


def root_motion(beat: float) -> tuple[float, float, float]:
    """Integrate the 18-count pattern through the supplied beat (0..72)."""
    beat = max(0.0, min(float(SHOW_BEATS), beat))
    whole = int(math.floor(beat))
    fraction = _smoothstep(beat - whole)
    x = y = yaw = 0.0

    for index in range(whole):
        dx, dy, dyaw = ROOT_STEPS[index % COUNTS_PER_WALL]
        world_dx, world_dy = _rotate(dx, dy, yaw)
        x += world_dx
        y += world_dy
        yaw += dyaw

    if whole < SHOW_BEATS:
        dx, dy, dyaw = ROOT_STEPS[whole % COUNTS_PER_WALL]
        world_dx, world_dy = _rotate(dx, dy, yaw)
        x += world_dx * fraction
        y += world_dy * fraction
        yaw += dyaw * fraction
    return x, y, yaw


def _quaternion(yaw: float, pitch: float, roll: float) -> tuple[float, float, float, float]:
    """MuJoCo wxyz quaternion from ZYX yaw, pitch, and roll."""
    cy, sy = math.cos(yaw / 2), math.sin(yaw / 2)
    cp, sp = math.cos(pitch / 2), math.sin(pitch / 2)
    cr, sr = math.cos(roll / 2), math.sin(roll / 2)
    return (
        cr * cp * cy + sr * sp * sy,
        sr * cp * cy - cr * sp * sy,
        cr * sp * cy + sr * cp * sy,
        cr * cp * sy - sr * sp * cy,
    )


def sample_pose(time_seconds: float, duck_index: int, base_xy: tuple[float, float]) -> DuckPose:
    """Sample one duck at show time, including subtle flock-specific attitude."""
    beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
    count = int(beat) % COUNTS_PER_WALL
    phase = beat - math.floor(beat)
    step_pulse = math.sin(math.pi * phase)
    groove = math.sin(2.0 * math.pi * (beat / 2.0 + (duck_index % 5) * 0.035))
    row = duck_index // 5
    col = duck_index % 5
    accent = (row - 2) * 0.018 + (col - 2) * 0.012

    travel_x, travel_y, yaw = root_motion(beat)
    z = 0.12 + 0.006 * abs(groove)
    pitch = 0.025 * groove
    roll = -0.035 * groove + accent

    pose = list(STAND_POSE)
    # Continuous knee compression gives the loose, grounded look in the clip.
    compression = 0.035 + 0.055 * abs(groove)
    pose[3] += compression
    pose[13] -= compression
    pose[2] += compression * 0.45
    pose[12] -= compression * 0.45

    # Alternating weight shift and foot pickup for traveling counts.
    if count in (0, 2, 4, 6, 8, 10):
        lead_left = count in (4, 6, 10)
        lift = 0.16 * step_pulse
        if lead_left:
            pose[2] += lift
            pose[3] += lift * 0.70
            pose[4] -= lift * 0.45
            roll += 0.055 * step_pulse
        else:
            pose[11] -= lift
            pose[12] -= lift * 0.70
            pose[13] += lift * 0.45
            roll -= 0.055 * step_pulse

    # Touches are visible: a toe point, a beak nod, and a tiny hold.
    if count in (3, 7, 11):
        pose[6] -= 0.13 * step_pulse
        pose[7] += 0.14 * groove
        z -= 0.008 * step_pulse

    # The reference clip's character section: low scoop forward, lean back,
    # then let the head lead the new wall.
    if count in (12, 13):
        scoop = step_pulse if count == 13 else _smoothstep(phase)
        z -= 0.022 * scoop
        pitch += 0.11 * scoop
        pose[5] += 0.24 * scoop
        pose[6] -= 0.22 * scoop
        pose[7] += (0.18 + accent) * scoop
        pose[8] -= 0.14 * groove
    elif count in (14, 15):
        lean = step_pulse if count == 15 else _smoothstep(phase)
        pitch -= 0.09 * lean
        pose[5] -= 0.18 * lean
        pose[6] += 0.15 * lean
        pose[8] += 0.18 * groove
    elif count in (16, 17):
        turn_lead = step_pulse if count == 16 else math.sin(math.pi * phase)
        pose[7] += 0.38 * turn_lead
        pose[8] += 0.13 * groove
        z += 0.012 * step_pulse

    # Keep the flock coordinated while the head-roll accent ripples by row.
    pose[8] += 0.075 * math.sin(2 * math.pi * (beat / 4.0 + row * 0.08))

    return DuckPose(
        position=(base_xy[0] + travel_x, base_xy[1] + travel_y, z),
        quaternion=_quaternion(yaw, pitch, roll),
        joints=tuple(pose),
    )


def sample_flock(time_seconds: float, grid: int = 5, spacing: float = 0.52) -> list[DuckPose]:
    positions = formation_positions(grid, spacing)
    return [sample_pose(time_seconds, index, xy) for index, xy in enumerate(positions)]

	"""Music-timed, 50 Hz Electric Slide choreography for a Microduck flock.

	The motion is deliberately kinematic for the first visual prototype: root poses
	and servo poses are sampled directly in MuJoCo. That makes the 25-duck show
	repeatable while keeping a clean seam for replacing the sampler with an ONNX
	policy later.
	"""

	from __future__ import annotations

	from dataclasses import dataclass
	import math


	BPM = 111.0
	BEAT_SECONDS = 60.0 / BPM
	COUNTS_PER_WALL = 18
	WALLS = 4
	SHOW_BEATS = COUNTS_PER_WALL * WALLS
	SHOW_SECONDS = SHOW_BEATS * BEAT_SECONDS
	CONTROL_HZ = 50

	# Official STAND keyframe, in official actuator/joint order.
	STAND_POSE = (
	0.0,
	-0.08726646259971647,
	-0.457924,
	-0.004940,
	0.452984,
	0.3490658503988659,
	0.3490658503988659,
	0.0,
	0.0,
	0.0,
	0.08726646259971647,
	0.457924,
	0.004940,
	-0.452984,
	)

	# Per-count root translation in the duck's local frame (forward, left) and yaw.
	# Right is negative local y. Four 18-count walls trace a closed square.
	ROOT_STEPS = (
	(0.00, -0.10, 0.0), # grapevine right: side
	(0.00, -0.04, 0.0), # cross behind
	(0.00, -0.10, 0.0), # side
	(0.00, 0.00, 0.0), # touch
	(0.00, 0.10, 0.0), # grapevine left: side
	(0.00, 0.04, 0.0), # cross behind
	(0.00, 0.10, 0.0), # side
	(0.00, 0.00, 0.0), # touch
	(-0.08, 0.00, 0.0), # walk back
	(-0.08, 0.00, 0.0),
	(-0.08, 0.00, 0.0),
	(0.00, 0.00, 0.0), # touch
	(0.08, 0.00, 0.0), # rock forward
	(0.00, 0.00, 0.0), # toe touch + scoop
	(-0.08, 0.00, 0.0), # rock back
	(0.00, 0.00, 0.0), # toe touch + lean
	(0.05, 0.00, 0.0), # step into turn
	(0.00, 0.00, math.pi / 2.0), # scuff + quarter turn left
	)


	@dataclass(frozen=True)
	class DuckPose:
	position: tuple[float, float, float]
	quaternion: tuple[float, float, float, float]
	joints: tuple[float, ...]


	def formation_positions(grid: int = 5, spacing: float = 0.52) -> list[tuple[float, float]]:
	"""Return a centered, row-major square formation."""
	center = (grid - 1) / 2.0
	return [
	((col - center) * spacing, (center - row) * spacing)
	for row in range(grid)
	for col in range(grid)
	]


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


	def _rotate(x: float, y: float, yaw: float) -> tuple[float, float]:
	cosine, sine = math.cos(yaw), math.sin(yaw)
	return cosine * x - sine * y, sine * x + cosine * y


	def root_motion(beat: float) -> tuple[float, float, float]:
	"""Integrate the 18-count pattern through the supplied beat (0..72)."""
	beat = max(0.0, min(float(SHOW_BEATS), beat))
	whole = int(math.floor(beat))
	fraction = _smoothstep(beat - whole)
	x = y = yaw = 0.0

	for index in range(whole):
	dx, dy, dyaw = ROOT_STEPS[index % COUNTS_PER_WALL]
	world_dx, world_dy = _rotate(dx, dy, yaw)
	x += world_dx
	y += world_dy
	yaw += dyaw

	if whole < SHOW_BEATS:
	dx, dy, dyaw = ROOT_STEPS[whole % COUNTS_PER_WALL]
	world_dx, world_dy = _rotate(dx, dy, yaw)
	x += world_dx * fraction
	y += world_dy * fraction
	yaw += dyaw * fraction
	return x, y, yaw


	def _quaternion(yaw: float, pitch: float, roll: float) -> tuple[float, float, float, float]:
	"""MuJoCo wxyz quaternion from ZYX yaw, pitch, and roll."""
	cy, sy = math.cos(yaw / 2), math.sin(yaw / 2)
	cp, sp = math.cos(pitch / 2), math.sin(pitch / 2)
	cr, sr = math.cos(roll / 2), math.sin(roll / 2)
	return (
	cr * cp * cy + sr * sp * sy,
	sr * cp * cy - cr * sp * sy,
	cr * sp * cy + sr * cp * sy,
	cr * cp * sy - sr * sp * cy,
	)


	def sample_pose(time_seconds: float, duck_index: int, base_xy: tuple[float, float]) -> DuckPose:
	"""Sample one duck at show time, including subtle flock-specific attitude."""
	beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
	count = int(beat) % COUNTS_PER_WALL
	phase = beat - math.floor(beat)
	step_pulse = math.sin(math.pi * phase)
	groove = math.sin(2.0 * math.pi * (beat / 2.0 + (duck_index % 5) * 0.035))
	row = duck_index // 5
	col = duck_index % 5
	accent = (row - 2) * 0.018 + (col - 2) * 0.012

	travel_x, travel_y, yaw = root_motion(beat)
	z = 0.12 + 0.006 * abs(groove)
	pitch = 0.025 * groove
	roll = -0.035 * groove + accent

	pose = list(STAND_POSE)
	# Continuous knee compression gives the loose, grounded look in the clip.
	compression = 0.035 + 0.055 * abs(groove)
	pose[3] += compression
	pose[13] -= compression
	pose[2] += compression * 0.45
	pose[12] -= compression * 0.45

	# Alternating weight shift and foot pickup for traveling counts.
	if count in (0, 2, 4, 6, 8, 10):
	lead_left = count in (4, 6, 10)
	lift = 0.16 * step_pulse
	if lead_left:
	pose[2] += lift
	pose[3] += lift * 0.70
	pose[4] -= lift * 0.45
	roll += 0.055 * step_pulse
	else:
	pose[11] -= lift
	pose[12] -= lift * 0.70
	pose[13] += lift * 0.45
	roll -= 0.055 * step_pulse

	# Touches are visible: a toe point, a beak nod, and a tiny hold.
	if count in (3, 7, 11):
	pose[6] -= 0.13 * step_pulse
	pose[7] += 0.14 * groove
	z -= 0.008 * step_pulse

	# The reference clip's character section: low scoop forward, lean back,
	# then let the head lead the new wall.
	if count in (12, 13):
	scoop = step_pulse if count == 13 else _smoothstep(phase)
	z -= 0.022 * scoop
	pitch += 0.11 * scoop
	pose[5] += 0.24 * scoop
	pose[6] -= 0.22 * scoop
	pose[7] += (0.18 + accent) * scoop
	pose[8] -= 0.14 * groove
	elif count in (14, 15):
	lean = step_pulse if count == 15 else _smoothstep(phase)
	pitch -= 0.09 * lean
	pose[5] -= 0.18 * lean
	pose[6] += 0.15 * lean
	pose[8] += 0.18 * groove
	elif count in (16, 17):
	turn_lead = step_pulse if count == 16 else math.sin(math.pi * phase)
	pose[7] += 0.38 * turn_lead
	pose[8] += 0.13 * groove
	z += 0.012 * step_pulse

	# Keep the flock coordinated while the head-roll accent ripples by row.
	pose[8] += 0.075 * math.sin(2 * math.pi * (beat / 4.0 + row * 0.08))

	return DuckPose(
	position=(base_xy[0] + travel_x, base_xy[1] + travel_y, z),
	quaternion=_quaternion(yaw, pitch, roll),
	joints=tuple(pose),
	)


	def sample_flock(time_seconds: float, grid: int = 5, spacing: float = 0.52) -> list[DuckPose]:
	positions = formation_positions(grid, spacing)
	return [sample_pose(time_seconds, index, xy) for index, xy in enumerate(positions)]