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