diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 95fa58f09..d1429d2b2 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -57,11 +57,35 @@ TOKEN_DIM = 64 # decoder latent size # added on top of this. DEFAULT_ANGLES = np.array( [ - -0.312, 0.0, 0.0, 0.669, -0.363, 0.0, - -0.312, 0.0, 0.0, 0.669, -0.363, 0.0, - 0.0, 0.0, 0.0, - 0.2, 0.2, 0.0, 0.6, 0.0, 0.0, 0.0, - 0.2, -0.2, 0.0, 0.6, 0.0, 0.0, 0.0, + -0.312, + 0.0, + 0.0, + 0.669, + -0.363, + 0.0, + -0.312, + 0.0, + 0.0, + 0.669, + -0.363, + 0.0, + 0.0, + 0.0, + 0.0, + 0.2, + 0.2, + 0.0, + 0.6, + 0.0, + 0.0, + 0.0, + 0.2, + -0.2, + 0.0, + 0.6, + 0.0, + 0.0, + 0.0, ], dtype=np.float32, ) @@ -110,10 +134,10 @@ def compute_kp_kd(): # Action-feature prefix for the latent-token interface (see _extract_token_from_action). -TOKEN_ACTION_PREFIX = "motion_token" +TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret # Proprio-state prefix for the token interface: the robot echoes the last commanded token # here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``. -TOKEN_STATE_PREFIX = "motion_token_state" +TOKEN_STATE_PREFIX = "motion_token_state" # nosec B105 - feature-key prefix, not a secret def token_action_key(i: int) -> str: @@ -142,10 +166,72 @@ INIT_RAMP_S = 3.0 # all-zero token, which is off the learned manifold and decodes to a slightly goofy stance. # This one decodes to a stable, natural standing pose. _NEUTRAL_TOKEN_CODES = np.array( - [-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1, - 2, -1, -3, -5, 3, 1, 1, -4, -1, -1, 1, -7, 0, 1, 2, -2, - 5, -2, -2, -4, 0, -1, 3, -1, 0, -5, -1, 0, -4, 0, 0, -1, - -1, 2, -2, 1, 3, 3, 1, 0, 0, 6, 0, -7, 3, 0, 2, -2], + [ + -1, + 3, + 1, + -1, + 1, + -3, + 6, + 1, + 1, + 1, + -2, + -4, + -2, + 0, + -3, + -1, + 2, + -1, + -3, + -5, + 3, + 1, + 1, + -4, + -1, + -1, + 1, + -7, + 0, + 1, + 2, + -2, + 5, + -2, + -2, + -4, + 0, + -1, + 3, + -1, + 0, + -5, + -1, + 0, + -4, + 0, + 0, + -1, + -1, + 2, + -2, + 1, + 3, + 3, + 1, + 0, + 0, + 6, + 0, + -7, + 3, + 0, + 2, + -2, + ], dtype=np.float32, ) NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0 # FSQ Div(16): integer codes -> on-grid token