Robot body orientations in body frame
(sim_data, motion_loader, t)
| 503 | |
| 504 | |
| 505 | def robot_body_ori_b(sim_data, motion_loader, t): |
| 506 | """Robot body orientations in body frame""" |
| 507 | # Get robot anchor pose |
| 508 | robot_anchor_pos = sim_data.body(motion_loader.anchor_body_name).xpos.copy().reshape(1, 3) |
| 509 | robot_anchor_quat = sim_data.body(motion_loader.anchor_body_name).xquat.copy().reshape(1, 4) |
| 510 | |
| 511 | # Get all robot body poses |
| 512 | robot_body_positions = [] |
| 513 | robot_body_orientations = [] |
| 514 | |
| 515 | for body_name in motion_loader.body_names: |
| 516 | body_pos = sim_data.body(body_name).xpos.copy().reshape(1, 3) |
| 517 | body_quat = sim_data.body(body_name).xquat.copy().reshape(1, 4) |
| 518 | robot_body_positions.append(body_pos) |
| 519 | robot_body_orientations.append(body_quat) |
| 520 | |
| 521 | robot_body_pos = np.concatenate(robot_body_positions, axis=0) # [14, 3] |
| 522 | robot_body_quat = np.concatenate(robot_body_orientations, axis=0) # [14, 4] |
| 523 | |
| 524 | # Transform to body frame - expand robot_anchor to match robot_body shape |
| 525 | robot_anchor_pos_expanded = robot_anchor_pos.repeat(len(motion_loader.body_names), axis=0) # [14, 3] |
| 526 | robot_anchor_quat_expanded = robot_anchor_quat.repeat(len(motion_loader.body_names), axis=0) # [14, 4] |
| 527 | |
| 528 | _, ori_b = subtract_frame_transforms( |
| 529 | torch.from_numpy(robot_anchor_pos_expanded).float(), |
| 530 | torch.from_numpy(robot_anchor_quat_expanded).float(), |
| 531 | torch.from_numpy(robot_body_pos).float(), |
| 532 | torch.from_numpy(robot_body_quat).float(), |
| 533 | ) |
| 534 | |
| 535 | # Convert to rotation matrix and take first 2 rows |
| 536 | mat = matrix_from_quat(ori_b) |
| 537 | mat_flat = mat[..., :2].reshape(-1) |
| 538 | |
| 539 | return np.array(mat_flat, dtype=np.float32).flatten() # [84] |
| 540 | |
| 541 | |
| 542 | def get_base_lin_vel(sim_data): |
no outgoing calls
no test coverage detected