File size: 6,798 Bytes
32502a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
"""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)]