Runs the simulation loop.
(sim: sim_utils.SimulationContext, scene: InteractiveScene, joint_names: list[str])
| 219 | |
| 220 | |
| 221 | def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene, joint_names: list[str]): |
| 222 | """Runs the simulation loop.""" |
| 223 | # Load motion |
| 224 | motion = MotionLoader( |
| 225 | motion_file=args_cli.input_file, |
| 226 | input_fps=args_cli.input_fps, |
| 227 | output_fps=args_cli.output_fps, |
| 228 | device=sim.device, |
| 229 | frame_range=args_cli.frame_range, |
| 230 | ) |
| 231 | |
| 232 | # Extract scene entities |
| 233 | robot = scene["robot"] |
| 234 | robot_joint_indexes = robot.find_joints(joint_names, preserve_order=True)[0] |
| 235 | |
| 236 | # ------- data logger ------------------------------------------------------- |
| 237 | log = { |
| 238 | "fps": [args_cli.output_fps], |
| 239 | "joint_pos": [], |
| 240 | "joint_vel": [], |
| 241 | "body_pos_w": [], |
| 242 | "body_quat_w": [], |
| 243 | "body_lin_vel_w": [], |
| 244 | "body_ang_vel_w": [], |
| 245 | "joint_names": np.array(robot.joint_names), |
| 246 | "body_names": np.array(robot.body_names), |
| 247 | } |
| 248 | file_saved = False |
| 249 | # -------------------------------------------------------------------------- |
| 250 | |
| 251 | # Simulation loop |
| 252 | while simulation_app.is_running(): |
| 253 | ( |
| 254 | ( |
| 255 | motion_base_pos, |
| 256 | motion_base_rot, |
| 257 | motion_base_lin_vel, |
| 258 | motion_base_ang_vel, |
| 259 | motion_dof_pos, |
| 260 | motion_dof_vel, |
| 261 | ), |
| 262 | reset_flag, |
| 263 | ) = motion.get_next_state() |
| 264 | |
| 265 | # set root state |
| 266 | root_states = robot.data.default_root_state.clone() |
| 267 | root_states[:, :3] = motion_base_pos |
| 268 | root_states[:, :2] += scene.env_origins[:, :2] |
| 269 | root_states[:, 3:7] = motion_base_rot |
| 270 | root_states[:, 7:10] = motion_base_lin_vel |
| 271 | root_states[:, 10:] = motion_base_ang_vel |
| 272 | robot.write_root_state_to_sim(root_states) |
| 273 | |
| 274 | # set joint state |
| 275 | joint_pos = robot.data.default_joint_pos.clone() |
| 276 | joint_vel = robot.data.default_joint_vel.clone() |
| 277 | joint_pos[:, robot_joint_indexes] = motion_dof_pos |
| 278 | joint_vel[:, robot_joint_indexes] = motion_dof_vel |
no test coverage detected