| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Literal |
|
|
| import numpy as np |
|
|
|
|
| CONTROL_HZ = 50 |
| OBSERVATION_HZ = 10 |
| CONTROL_DECIMATION = CONTROL_HZ // OBSERVATION_HZ |
| |
| CHUNK_SIZE = 100 |
| EPISODE_DURATION_S = 5.0 |
| TERMINAL_STABILIZATION_S = 1.5 |
| RECOVERY_STABLE_S = 0.5 |
| STATE_DIM = 61 |
| ACTION_DIM = 14 |
| EMBODIMENT_ID = 24 |
|
|
| POLICY_REVISION = "916c1e1ee11f8460506532e037a4cf1aa110aca6" |
| MICRODUCK_RL_REVISION = "d424a0c899f6b33cbd3daeb279913134349c0b63" |
| TINYVLA_REVISION = "565d139f4960174fdd3006ec57c25ac62e129245" |
|
|
| DEFAULT_POSE = np.array( |
| [ |
| 0.0, |
| -0.0873, |
| -0.4579, |
| -0.0049, |
| 0.4530, |
| 0.3491, |
| 0.3491, |
| 0.0, |
| 0.0, |
| 0.0, |
| 0.0873, |
| 0.4579, |
| 0.0049, |
| -0.4530, |
| ], |
| dtype=np.float32, |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class BehaviorSpec: |
| name: str |
| policy_file: str |
| scene_file: str |
| duration_s: float |
| start_keyframe: str |
| handoff_policy_file: str | None = None |
| handoff_after_s: float | None = None |
| phase_period_s: float | None = None |
| kind: Literal[ |
| "stand", |
| "walking", |
| "sitstand", |
| "phase", |
| "kick", |
| "roulade", |
| "roller", |
| "recovery", |
| ] = "stand" |
|
|
|
|
| BEHAVIORS: tuple[BehaviorSpec, ...] = ( |
| BehaviorSpec( |
| "stand", |
| "alpha_stand.onnx", |
| "scene.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| kind="stand", |
| ), |
| BehaviorSpec( |
| "walking", |
| "alpha_walking.onnx", |
| "scene.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| handoff_policy_file="alpha_stand.onnx", |
| handoff_after_s=4.0, |
| kind="walking", |
| ), |
| BehaviorSpec( |
| "sitstand", |
| "alpha_sitstand.onnx", |
| "scene.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| kind="sitstand", |
| ), |
| BehaviorSpec( |
| "ground_pick", |
| "alpha_ground_pick.onnx", |
| "scene.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| handoff_policy_file="alpha_stand.onnx", |
| handoff_after_s=2.8, |
| phase_period_s=4.0, |
| kind="phase", |
| ), |
| BehaviorSpec( |
| "kick_left", |
| "ball_kick_left.onnx", |
| "scene_ball.xml", |
| 5.0, |
| "STAND", |
| handoff_policy_file="alpha_stand.onnx", |
| handoff_after_s=3.0, |
| kind="kick", |
| ), |
| BehaviorSpec( |
| "kick_right", |
| "ball_kick_right.onnx", |
| "scene_ball.xml", |
| 5.0, |
| "STAND", |
| handoff_policy_file="alpha_stand.onnx", |
| handoff_after_s=3.0, |
| kind="kick", |
| ), |
| BehaviorSpec( |
| "roller", |
| "roller.onnx", |
| "scene_rollers.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| handoff_policy_file="alpha_stand.onnx", |
| handoff_after_s=4.0, |
| kind="roller", |
| ), |
| BehaviorSpec( |
| "roller_crouch", |
| "roller_crouch.onnx", |
| "scene_rollers.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| handoff_policy_file="roller.onnx", |
| handoff_after_s=3.5, |
| phase_period_s=5.0, |
| kind="phase", |
| ), |
| BehaviorSpec( |
| "roulade", |
| "roulade.onnx", |
| "scene.xml", |
| 5.0, |
| "STAND", |
| handoff_policy_file="alpha_stand.onnx", |
| handoff_after_s=2.0, |
| kind="roulade", |
| ), |
| ) |
|
|
|
|
| RECOVERY_BEHAVIORS: tuple[BehaviorSpec, ...] = ( |
| BehaviorSpec( |
| "recovery", |
| "alpha_stand.onnx", |
| "scene.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| kind="recovery", |
| ), |
| BehaviorSpec( |
| "recovery_roller", |
| "alpha_stand.onnx", |
| "scene_rollers.xml", |
| EPISODE_DURATION_S, |
| "STAND", |
| kind="recovery", |
| ), |
| ) |
|
|
| ALL_BEHAVIORS: tuple[BehaviorSpec, ...] = BEHAVIORS + RECOVERY_BEHAVIORS |
|
|
|
|
| def split_for_ordinal(ordinal: int, episodes_per_policy: int) -> str: |
| """Put the final 10% of each teacher into validation.""" |
| n_validation = max(1, round(episodes_per_policy * 0.1)) |
| return "validation" if ordinal >= episodes_per_policy - n_validation else "train" |
|
|
|
|
| def validate_contract() -> None: |
| if CONTROL_HZ % OBSERVATION_HZ: |
| raise ValueError("observation rate must divide the control rate") |
| if len(DEFAULT_POSE) != ACTION_DIM: |
| raise ValueError("default pose and action dimensions disagree") |
| if any(b.duration_s != EPISODE_DURATION_S for b in ALL_BEHAVIORS): |
| raise ValueError("every episode must be exactly five seconds") |
| names = [b.name for b in ALL_BEHAVIORS] |
| if len(names) != len(set(names)): |
| raise ValueError("behavior names must be unique") |
|
|
