Compute the observation vector from current state
(d, config, action, cmd, height_cmd, n_joints)
| 57 | return quat_rotate_inverse(quat, gravity_vec) |
| 58 | |
| 59 | def compute_observation(d, config, action, cmd, height_cmd, n_joints): |
| 60 | """Compute the observation vector from current state""" |
| 61 | # Get state from MuJoCo |
| 62 | qj = d.qpos[7:7+n_joints].copy() |
| 63 | dqj = d.qvel[6:6+n_joints].copy() |
| 64 | quat = d.qpos[3:7].copy() |
| 65 | omega = d.qvel[3:6].copy() |
| 66 | |
| 67 | # Handle default angles padding |
| 68 | if len(config['default_angles']) < n_joints: |
| 69 | padded_defaults = np.zeros(n_joints, dtype=np.float32) |
| 70 | padded_defaults[:len(config['default_angles'])] = config['default_angles'] |
| 71 | else: |
| 72 | padded_defaults = config['default_angles'][:n_joints] |
| 73 | |
| 74 | # Scale the values |
| 75 | qj_scaled = (qj - padded_defaults) * config['dof_pos_scale'] |
| 76 | dqj_scaled = dqj * config['dof_vel_scale'] |
| 77 | gravity_orientation = get_gravity_orientation(quat) |
| 78 | omega_scaled = omega * config['ang_vel_scale'] |
| 79 | |
| 80 | # Calculate single observation dimension |
| 81 | single_obs_dim = 3 + 1 + 3 + 3 + n_joints + n_joints + 12 |
| 82 | |
| 83 | # Create single observation |
| 84 | single_obs = np.zeros(single_obs_dim, dtype=np.float32) |
| 85 | single_obs[0:3] = cmd[:3] * config['cmd_scale'] |
| 86 | single_obs[3:4] = np.array([height_cmd]) |
| 87 | single_obs[4:7] = omega_scaled |
| 88 | single_obs[7:10] = gravity_orientation |
| 89 | single_obs[10:10+n_joints] = qj_scaled |
| 90 | single_obs[10+n_joints:10+2*n_joints] = dqj_scaled |
| 91 | single_obs[10+2*n_joints:10+2*n_joints+12] = action |
| 92 | |
| 93 | return single_obs, single_obs_dim |
| 94 | |
| 95 | def main(): |
| 96 | # Load configuration |
no test coverage detected