File size: 2,951 Bytes
d8290eb | 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 | """Policy observation assembly (``PolicyCfg`` of the Online-CLoSD ball task).
Order and semantics (N = number of actuated joints):
task_command 9 | ref joint_pos N | ref joint_vel N | root xy 2 | ref pelvis xy 2 |
ref pelvis 6-D 6 | projected gravity 3 | heading xy 2 | base ang vel (body) 3 |
joint_pos - default N | joint_vel N | previous training-clipped action N -> 27 + 5N
Noise is not added (deployment = clean sensors); enable ``noise=True`` for a
training-like corruption.
"""
from __future__ import annotations
import numpy as np
from .mathutil import quat_apply, quat_apply_inverse, rot6d_interleaved
_TASK_NOISE = np.array([0.005, 0.005, 0.008, 0.008, 0.05, 0.05, 0.05, 0.03, 0.03])
def build_observation(
*,
task_command: np.ndarray,
ref_joint_pos: np.ndarray,
ref_joint_vel: np.ndarray,
root_xy: np.ndarray,
ref_pelvis_xy: np.ndarray,
ref_pelvis_quat: np.ndarray,
pelvis_quat: np.ndarray,
heading_quat: np.ndarray | None,
base_ang_vel_body: np.ndarray,
joint_pos: np.ndarray,
joint_vel: np.ndarray,
default_joint_pos: np.ndarray,
prev_action: np.ndarray,
noise: bool = False,
rng: np.random.Generator | None = None,
) -> np.ndarray:
gravity = quat_apply_inverse(pelvis_quat, np.array([0.0, 0.0, -1.0]))
# Gravity and angular velocity come from the robot IMU. Heading can use
# the sign-normalised mocap quaternion to avoid unobservable IMU yaw drift.
heading_source = pelvis_quat if heading_quat is None else heading_quat
heading = quat_apply(heading_source, np.array([1.0, 0.0, 0.0]))[:2]
heading = heading / max(float(np.linalg.norm(heading)), 1e-6)
parts = [
np.asarray(task_command, dtype=np.float64),
np.asarray(ref_joint_pos, dtype=np.float64),
np.asarray(ref_joint_vel, dtype=np.float64),
np.asarray(root_xy, dtype=np.float64),
np.asarray(ref_pelvis_xy, dtype=np.float64),
rot6d_interleaved(np.asarray(ref_pelvis_quat, dtype=np.float64)),
gravity,
heading,
np.asarray(base_ang_vel_body, dtype=np.float64),
np.asarray(joint_pos, dtype=np.float64) - np.asarray(default_joint_pos, dtype=np.float64),
np.asarray(joint_vel, dtype=np.float64),
np.asarray(prev_action, dtype=np.float64),
]
if noise:
rng = rng or np.random.default_rng()
n = len(joint_pos)
parts[0] = parts[0] + rng.uniform(-_TASK_NOISE, _TASK_NOISE)
parts[3] = parts[3] + rng.uniform(-0.05, 0.05, 2)
parts[6] = parts[6] + rng.uniform(-0.03, 0.03, 3)
parts[7] = parts[7] + rng.uniform(-0.04, 0.04, 2)
parts[8] = parts[8] + rng.uniform(-0.2, 0.2, 3)
parts[9] = parts[9] + rng.uniform(-0.01, 0.01, n)
parts[10] = parts[10] + rng.uniform(-1.0, 1.0, n)
return np.concatenate(parts).astype(np.float32)
def expected_obs_dim(num_actions: int) -> int:
return 27 + 5 * num_actions
|