| 22 | |
| 23 | |
| 24 | class AnalysisPlotForce(RL_EvalCallback): |
| 25 | training_loop: PPO |
| 26 | env: LeggedRobotBase |
| 27 | |
| 28 | def __init__(self, config, training_loop: PPO): |
| 29 | super().__init__(config, training_loop) |
| 30 | env: LeggedRobotBase = self.training_loop.env |
| 31 | self.env = env |
| 32 | self.policy = self.training_loop._get_inference_policy() |
| 33 | self.num_envs = self.env.num_envs |
| 34 | self.logger = WebLogger(self.config.sim_dt) |
| 35 | self.reset_buffers() |
| 36 | self.log_single_robot = self.config.log_single_robot |
| 37 | |
| 38 | def reset_buffers(self): |
| 39 | self.obs_buf = [[] for _ in range(self.num_envs)] |
| 40 | self.critic_obs_buf = [[] for _ in range(self.num_envs)] |
| 41 | self.act_buf = [[] for _ in range(self.num_envs)] |
| 42 | |
| 43 | def on_pre_evaluate_policy(self): |
| 44 | # Doing this in two lines because of type annotation issues. |
| 45 | self.robot_num_dofs = self.env.num_dofs |
| 46 | self.log_dof_pos_limits = self.env.dof_pos_limits.cpu().numpy() |
| 47 | self.log_dof_vel_limits = self.env.dof_vel_limits.cpu().numpy() |
| 48 | self.log_dof_torque_limits = self.env.torque_limits.cpu().numpy() |
| 49 | self.logger.set_robot_limits(self.log_dof_pos_limits, self.log_dof_vel_limits, self.log_dof_torque_limits) |
| 50 | self.logger.set_robot_num_dofs(self.robot_num_dofs) |
| 51 | |
| 52 | def on_post_evaluate_policy(self): |
| 53 | pass |
| 54 | |
| 55 | def on_pre_eval_env_step(self, actor_state): |
| 56 | obs: Tensor = actor_state["obs"]["actor_obs"].cpu() |
| 57 | critic_obs: Tensor = actor_state["obs"]["critic_obs"].cpu() |
| 58 | actions: Tensor = actor_state["actions"].cpu() |
| 59 | |
| 60 | for i in range(self.num_envs): |
| 61 | self.obs_buf[i].append(obs[i]) |
| 62 | self.critic_obs_buf[i].append(critic_obs[i]) |
| 63 | self.act_buf[i].append(actions[i]) |
| 64 | |
| 65 | rigid_body_pos = self.env._rigid_body_pos_extend |
| 66 | ref_body_pos = self.env.ref_body_pos_extend |
| 67 | vr_id = self.env.motion_tracking_id |
| 68 | |
| 69 | |
| 70 | |
| 71 | if self.log_single_robot: |
| 72 | self.logger.log_states( |
| 73 | { |
| 74 | 'dof_pos_target': actions[0].cpu().numpy(), |
| 75 | 'dof_pos': self.env.simulator.dof_pos[0].cpu().numpy(), |
| 76 | 'dof_vel': self.env.simulator.dof_vel[0].cpu().numpy(), |
| 77 | 'dof_torque': self.env.torques[0].cpu().numpy(), |
| 78 | 'base_vel_x': self.env.base_lin_vel[0, 0].item(), |
| 79 | 'base_vel_y': self.env.base_lin_vel[0, 1].item(), |
| 80 | 'base_vel_z': self.env.base_lin_vel[0, 2].item(), |
| 81 | 'base_vel_yaw': self.env.base_ang_vel[0, 2].item(), |
nothing calls this directly
no outgoing calls
no test coverage detected