Advance one control step using the current policy.
(self)
| 179 | time.sleep(self.config.control_dt) |
| 180 | |
| 181 | def run(self): |
| 182 | """Advance one control step using the current policy.""" |
| 183 | self.counter += 1 |
| 184 | # Get the current joint position and velocity |
| 185 | for default_idx, sorted_idx in self.config.body_default_sorted_idx_tuples: |
| 186 | self.qj[sorted_idx] = self.low_state.motor_state[default_idx].q |
| 187 | self.dqj[sorted_idx] = self.low_state.motor_state[default_idx].dq |
| 188 | |
| 189 | # imu_state quaternion: w, x, y, z |
| 190 | quat = self.low_state.imu_state.quaternion |
| 191 | # angular velocity handling |
| 192 | # for torso IMU, transform requires shape (1, 3); otherwise keep (3,) |
| 193 | if self.config.imu_type == "torso": |
| 194 | ang_vel_in = np.array([self.low_state.imu_state.gyroscope], dtype=np.float32) |
| 195 | # h1 and h1_2 imu is on the torso |
| 196 | # imu data needs to be transformed to the pelvis frame |
| 197 | waist_yaw = self.low_state.motor_state[self.config.arm_waist_joint2motor_idx[0]].q |
| 198 | waist_yaw_omega = self.low_state.motor_state[self.config.arm_waist_joint2motor_idx[0]].dq |
| 199 | quat, ang_vel = transform_imu_data( |
| 200 | waist_yaw=waist_yaw, waist_yaw_omega=waist_yaw_omega, imu_quat=quat, imu_omega=ang_vel_in |
| 201 | ) |
| 202 | else: |
| 203 | ang_vel = np.asarray(self.low_state.imu_state.gyroscope, dtype=np.float32) |
| 204 | |
| 205 | # create observation |
| 206 | gravity_orientation = get_gravity_orientation(quat) |
| 207 | base_euler_xyz = get_euler_xyz(quat) |
| 208 | qj_obs = self.qj.copy() |
| 209 | dqj_obs = self.dqj.copy() |
| 210 | qj_obs = (qj_obs - self.config.default_angles) * self.config.dof_pos_scale |
| 211 | dqj_obs = dqj_obs * self.config.dof_vel_scale |
| 212 | ang_vel = ang_vel * self.config.ang_vel_scale |
| 213 | period = 0.8 |
| 214 | count = self.counter * self.config.control_dt |
| 215 | phase = count % period / period |
| 216 | sin_phase = np.sin(2 * np.pi * phase) |
| 217 | cos_phase = np.cos(2 * np.pi * phase) |
| 218 | |
| 219 | self.cmd[0] = self.remote_controller.ly |
| 220 | self.cmd[1] = self.remote_controller.lx * -1 |
| 221 | self.cmd[2] = self.remote_controller.rx * -1 |
| 222 | |
| 223 | num_actions = self.config.num_actions |
| 224 | self.obs[0:3] = self.cmd * self.config.cmd_scale |
| 225 | self.obs[3:6] = ang_vel |
| 226 | self.obs[6:9] = gravity_orientation |
| 227 | self.obs[9 : 9 + num_actions] = qj_obs |
| 228 | self.obs[9 + num_actions : 9 + num_actions * 2] = dqj_obs |
| 229 | self.obs[9 + num_actions * 2 : 9 + num_actions * 3] = self.action |
| 230 | self.obs_queue.append(deepcopy(self.obs)) |
| 231 | |
| 232 | # Get the action from the policy network |
| 233 | obs_tensor = torch.from_numpy(np.concatenate(list(self.obs_queue), axis=0)).unsqueeze(0) |
| 234 | self.action = self.policy(obs_tensor).detach().numpy().squeeze() |
| 235 | delay = 0.2 |
| 236 | self.action = (1 - delay) * self.action + delay * self.last_action |
| 237 | self.last_action = self.action.copy() |
| 238 |