| """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])) |
| |
| |
| 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 |
|
|
