from __future__ import annotations import math from pathlib import Path import sys import tempfile import unittest import xml.etree.ElementTree as ET DANCE_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(DANCE_DIR)) from choreography import ( # noqa: E402 COUNTS_PER_WALL, SHOW_BEATS, SHOW_SECONDS, formation_positions, root_motion, sample_flock, ) from generate_scene import START_QUATERNION, build_scene # noqa: E402 from physics_choreography import ( # noqa: E402 MOTION_TIMELINE, command_for_time, policy_for_time, ) from literal_choreography import ( # noqa: E402 LITERAL_COUNTS, command_vector as literal_command_vector, joint_bias as literal_joint_bias, policy_for_time as literal_policy_for_time, ) from simple_choreography import ( # noqa: E402 COUNTS_PER_WALL as SIMPLE_COUNTS_PER_WALL, SHOW_SECONDS as SIMPLE_SHOW_SECONDS, SIMPLE_COUNTS, command_vector as simple_command_vector, joint_bias as simple_joint_bias, policy_for_time as simple_policy_for_time, ) class ChoreographyTests(unittest.TestCase): def test_simple_sequence_matches_supplied_16_counts(self) -> None: self.assertEqual(len(SIMPLE_COUNTS), 16) self.assertEqual( [item.label for item in SIMPLE_COUNTS], [ "right step right", "left follows right", "right step right", "left tap + beak-snap head bob", "left step left", "right follows left", "left step left", "right tap + beak-snap head bob", "right step back", "left step back", "right step back", "left tap", "rock forward", "rock back", "left step forward", "right scuff + pivot 90 degrees left", ], ) def test_simple_controller_keeps_one_continuous_balance_policy(self) -> None: samples = { simple_policy_for_time(index / 50.0) for index in range(int(SIMPLE_SHOW_SECONDS * 50)) } self.assertEqual(samples, {"walk"}) def test_simple_targets_fit_policy_and_joint_contracts(self) -> None: beat_seconds = SIMPLE_SHOW_SECONDS / (SIMPLE_COUNTS_PER_WALL * 4) for count in range(SIMPLE_COUNTS_PER_WALL): time_seconds = (count + 0.5) * beat_seconds command = simple_command_vector(time_seconds) bias = simple_joint_bias(time_seconds) self.assertEqual(command.shape, (13,)) self.assertEqual(bias.shape, (14,)) self.assertTrue(all(math.isfinite(float(value)) for value in command)) self.assertTrue(all(math.isfinite(float(value)) for value in bias)) if count == 15: self.assertGreater(command[2], 0.0) else: self.assertEqual(command[2], 0.0) def test_literal_sequence_matches_canonical_18_counts(self) -> None: self.assertEqual(len(LITERAL_COUNTS), 18) self.assertEqual( [item.label for item in LITERAL_COUNTS], [ "step right", "left crosses behind", "step right", "left touch/stomp", "step left", "right crosses behind", "step left", "right touch/stomp", "right step back", "left step back", "right step back", "left touch", "left step forward", "right touch", "right step back", "left touch", "left step forward", "right scuff + quarter-turn left", ], ) def test_literal_controller_keeps_one_continuous_balance_policy(self) -> None: samples = { literal_policy_for_time(index / 50.0) for index in range(int(SHOW_SECONDS * 50)) } self.assertEqual(samples, {"walk"}) def test_literal_targets_fit_policy_and_joint_contracts(self) -> None: for count in range(COUNTS_PER_WALL): time_seconds = (count + 0.5) * (SHOW_SECONDS / SHOW_BEATS) command = literal_command_vector(time_seconds) bias = literal_joint_bias(time_seconds) self.assertEqual(command.shape, (13,)) self.assertEqual(bias.shape, (14,)) self.assertTrue(all(math.isfinite(float(value)) for value in command)) self.assertTrue(all(math.isfinite(float(value)) for value in bias)) if count == 17: self.assertGreater(command[2], 0.0) else: self.assertEqual(command[2], 0.0) def test_formation_is_centered_and_collision_safe(self) -> None: points = formation_positions() self.assertEqual(len(points), 25) self.assertAlmostEqual(sum(x for x, _ in points), 0.0) self.assertAlmostEqual(sum(y for _, y in points), 0.0) minimum = min( math.dist(a, b) for index, a in enumerate(points) for b in points[index + 1 :] ) self.assertGreaterEqual(minimum, 0.50) def test_four_walls_close_the_root_path(self) -> None: x, y, yaw = root_motion(SHOW_BEATS) self.assertAlmostEqual(x, 0.0, places=8) self.assertAlmostEqual(y, 0.0, places=8) self.assertAlmostEqual(yaw, 2.0 * math.pi, places=8) def test_each_wall_turns_exactly_one_quarter(self) -> None: for wall in range(1, 5): _, _, yaw = root_motion(wall * COUNTS_PER_WALL) self.assertAlmostEqual(yaw, wall * math.pi / 2.0, places=8) def test_flock_pose_dimensions(self) -> None: flock = sample_flock(SHOW_SECONDS * 0.37) self.assertEqual(len(flock), 25) for pose in flock: self.assertEqual(len(pose.quaternion), 4) self.assertEqual(len(pose.joints), 14) def test_physics_commands_fit_official_policy_contract(self) -> None: for duck in range(25): command = command_for_time(SHOW_SECONDS * 0.41, duck) self.assertEqual(len(command.twist), 3) self.assertEqual(len(command.head_pose), 4) self.assertEqual(len(command.body_pose), 6) self.assertLessEqual(abs(command.twist[0]), 0.3) self.assertLessEqual(abs(command.twist[1]), 0.3) self.assertEqual(command.twist[1], 0.0) self.assertGreaterEqual(command.twist[0], 0.0) self.assertLessEqual(abs(command.twist[2]), 1.0) def test_motion_derived_accents_use_both_one_shot_policies(self) -> None: self.assertTrue(MOTION_TIMELINE.exists()) samples = { policy_for_time(index / 50.0) for index in range(int(SHOW_SECONDS * 50)) } self.assertEqual(samples, {"walk", "kick_left", "kick_right"}) def test_motion_derived_commands_remain_in_deployment_ranges(self) -> None: for index in range(int(SHOW_SECONDS * 50)): command = command_for_time(index / 50.0) self.assertEqual(command.twist[1], 0.0) self.assertGreaterEqual(command.twist[0], 0.0) self.assertLessEqual(command.twist[0], 0.3) self.assertLessEqual(abs(command.twist[2]), 1.0) self.assertLessEqual(max(abs(value) for value in command.head_pose), 0.20) self.assertLessEqual(max(abs(value) for value in command.body_pose), 0.12) class SceneGeneratorTests(unittest.TestCase): def test_scene_has_unique_per_duck_controls(self) -> None: with tempfile.TemporaryDirectory() as directory: output = Path(directory) / "scene.xml" build_scene(output) root = ET.parse(output).getroot() self.assertEqual(root.find("option").get("timestep"), "0.005") trunks = [ body for body in root.findall(".//body") if body.get("name", "").endswith("trunk_base") ] actuator_section = root.find("actuator") sensor_section = root.find("sensor") self.assertIsNotNone(actuator_section) self.assertIsNotNone(sensor_section) actuators = list(actuator_section) sensors = list(sensor_section) self.assertEqual(len(trunks), 25) self.assertEqual(len(actuators), 350) self.assertEqual(len(sensors), 150) self.assertEqual(len({item.get("name") for item in actuators}), 350) self.assertEqual(len({item.get("name") for item in sensors}), 150) self.assertEqual(actuators[0].get("joint"), "duck_00_left_hip_yaw") self.assertEqual(actuators[-1].get("joint"), "duck_24_right_ankle") key = root.find("./keyframe/key") qpos = [float(value) for value in key.get("qpos").split()] for actual, expected in zip(qpos[3:7], START_QUATERNION, strict=True): self.assertAlmostEqual(actual, expected, places=8) def test_scene_uses_light_sky_and_checkerboard_floor(self) -> None: with tempfile.TemporaryDirectory() as directory: output = Path(directory) / "scene.xml" build_scene(output) root = ET.parse(output).getroot() textures = {item.get("name"): item for item in root.findall("./asset/texture")} checker = textures["dance_floor"] self.assertEqual(checker.get("builtin"), "checker") tiles = [ geom for geom in root.findall("./worldbody/geom") if geom.get("name", "").startswith("checker_") ] self.assertEqual(len(tiles), 196) self.assertEqual({tile.get("material") for tile in tiles}, {"tile_light", "tile_slate"}) sky = next( item for item in root.findall("./asset/texture") if item.get("type") == "skybox" ) self.assertGreater(float(sky.get("rgb1", "0").split()[0]), 0.9) if __name__ == "__main__": unittest.main()