(self)
| 177 | viewport.width // self.macro_block_size * self.macro_block_size) |
| 178 | macro_block_size = 16 |
| 179 | def render_step(self): |
| 180 | if self.is_render: |
| 181 | # if self.viewer.is_running(): |
| 182 | if self.viewer.is_alive: |
| 183 | self.viewer.render() |
| 184 | # self.viewer.sync() |
| 185 | if self.is_recording and (time.time() - self.start_time >= 1/self.fps): |
| 186 | self.start_time = time.time() |
| 187 | mujoco.mjr_readPixels(self._frame_buffer, None, self.viewer.viewport, self.viewer.ctx) |
| 188 | # imageio.imwrite('frame.png', frame[::-1]) |
| 189 | # breakpoint() |
| 190 | |
| 191 | self._video_buffer.append(self._frame_buffer[::-1].copy()[:self._frame_size[0], :self._frame_size[1]]) |
| 192 | else: |
| 193 | if self.is_recording: |
| 194 | logger.info("Mujoco: Saving video ...") |
| 195 | |
| 196 | model_path:Path = self.cfg.checkpoint |
| 197 | model_id = model_path.stem.replace('model_', 'ckpt_') |
| 198 | video_path:Path = model_path.parent.parent / 'renderings' / model_id / f'video_{time.strftime("%Y%m%d_%H%M%S")}.mp4' |
| 199 | video_path.parent.mkdir(parents=True, exist_ok=True) |
| 200 | import imageio |
| 201 | imageio.mimsave(video_path, self._video_buffer, fps=self.fps, macro_block_size=self.macro_block_size) |
| 202 | logger.info(f"Video saved to {video_path}") |
| 203 | raise RobotExitException("Mujoco Robot Exit") |
| 204 | |
| 205 | |
| 206 | class MujocoRobot(URCIRobot, ViewerPlugin): |
no test coverage detected