Repository navigation
Add autogeneration of doc media - #146
Merged
Merged
Conversation
amacati
requested changes
Oct 28, 2026
Comment on lines
+141
to
+187
| def render( | ||
| sim: Sim, | ||
| mode: str | None = "human", | ||
| world: int = 0, | ||
| camera: int | str = -1, | ||
| cam_config: dict | None = None, | ||
| width: int = 1920, | ||
| height: int = 1080, | ||
| ) -> np.ndarray | None: | ||
| if mode != "human": | ||
| return sim_render(sim, mode, world, camera, cam_config, width, height) | ||
| if "video" not in kinds: | ||
| return None | ||
| if docs_camera is not None: | ||
| camera = docs_camera | ||
| elif docs_cam_config: | ||
| camera = -1 | ||
| cam_config = CAM_CONFIG | (cam_config or {}) | (docs_cam_config or {}) | ||
| if sim.viewer is None: | ||
| videos.append(Video(out / f"video-{len(videos) + 1:02d}.mp4")) | ||
| if docs_geomgroup is not None: | ||
| sim_render(sim, "rgb_array", world, camera, cam_config, WIDTH, HEIGHT) | ||
| sim.viewer.viewer.vopt.geomgroup[:] = [ | ||
| group in docs_geomgroup for group in range(mujoco.mjNGROUP) | ||
| ] | ||
| rendered = videos[-1].add( | ||
| float(sim.time[world, 0]), | ||
| lambda: sim_render(sim, "rgb_array", world, camera, cam_config, WIDTH, HEIGHT), | ||
| ) | ||
| if not rendered: # Only a render clears the markers, so drop those of skipped frames | ||
| sim.viewer.viewer._markers.clear() | ||
|
|
||
| def func_animation(self: animation.FuncAnimation, fig: plt.Figure, *args: Any, **kwargs: Any): | ||
| func_animation_init(self, fig, *args, **kwargs) | ||
| animations[fig] = self | ||
|
|
||
| def show(*args: Any, **kwargs: Any): | ||
| if "plot" in kinds: | ||
| for num in plt.get_fignums(): | ||
| fig = plt.figure(num) | ||
| save_figure(fig, out / f"plot-{next(plot_ids):02d}", animations.get(fig)) | ||
| plt.close("all") | ||
|
|
||
| sys.argv = [str(example), *(["--plot"] if "plot" in kinds else [])] | ||
| spec = importlib.util.spec_from_file_location("__main__", example) | ||
| module = importlib.util.module_from_spec(spec) | ||
| sys.modules["__main__"] = module |
Collaborator
There was a problem hiding this comment.
I get what you are trying to do, and I'm not saying this can easily be done otherwise, but this is some voodoo-level stuff. We could also maybe render as RGB in the examples, that would mean one more argument that is always "human" by default, and then you only have to wrap sim.render() with a minimal wrapper.
Comment on lines
+36
to
+93
| class Video: | ||
| """H.264 video of rendered frames, resampled to real time from the sim time of each frame. | ||
|
|
||
| A video of a single frame is saved as png. | ||
| """ | ||
|
|
||
| def __init__(self, path: Path): | ||
| """Create the video, the file is written on the second frame.""" | ||
| self.path = path | ||
| self.writer = None | ||
| self.first = None | ||
| self.n_frames = 0 | ||
| self.start = None | ||
| self.offset = 0.0 | ||
| self.last = 0.0 | ||
|
|
||
| def add(self, t: float, render: Callable[[], np.ndarray]) -> bool: | ||
| """Render and write a frame at sim time ``t`` if the video needs one.""" | ||
| if t + self.offset < self.last: # The sim was reset | ||
| self.offset = self.last + 1 / FPS - t | ||
| t = self.last = t + self.offset | ||
| if self.start is None: | ||
| self.start = t | ||
| if t < self._next(): | ||
| return False | ||
| frame = np.ascontiguousarray(render()) | ||
| while self._next() <= t: | ||
| self._write(frame) | ||
| self.n_frames += 1 | ||
| return True | ||
|
|
||
| def close(self): | ||
| """Finish writing the video.""" | ||
| if self.writer is not None: | ||
| self.writer.close() | ||
| elif self.first is not None: | ||
| iio.imwrite(self.path.with_suffix(".png"), self.first) | ||
|
|
||
| def _write(self, frame: np.ndarray): | ||
| if self.first is None: | ||
| self.first = frame | ||
| return | ||
| if self.writer is None: | ||
| self.writer = imageio_ffmpeg.write_frames( | ||
| str(self.path), | ||
| (frame.shape[1], frame.shape[0]), | ||
| fps=FPS, | ||
| codec="libx264", | ||
| quality=None, | ||
| macro_block_size=2, | ||
| output_params=["-crf", CRF, "-preset", "slow", "-movflags", "+faststart"], | ||
| ) | ||
| self.writer.send(None) | ||
| self.writer.send(self.first) | ||
| self.writer.send(frame) | ||
|
|
||
| def _next(self) -> float: | ||
| return self.start + self.n_frames / FPS - 1e-9 |
Collaborator
There was a problem hiding this comment.
Do we have to do this on our own? There is nothing in PIL or another Python package that we can use?
amacati
approved these changes
Oct 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
We previously has some selected media in the docs for some of the examples, but
This PR let's you define what output of the example to automatically capture and put into the docs, i.e., console output, renders, and plots can be automatically captured and added in the correct location.
Closes #140