fix: replace moviepy with PyAV in animate() to prevent last frame drop#2704
Closed
vlordier wants to merge 14 commits intoGenesis-Embodied-AI:mainfrom
Closed
fix: replace moviepy with PyAV in animate() to prevent last frame drop#2704vlordier wants to merge 14 commits intoGenesis-Embodied-AI:mainfrom
vlordier wants to merge 14 commits intoGenesis-Embodied-AI:mainfrom
Conversation
- Add history_length parameter to ContactForce sensor options - Override read() to return historical force readings from ring buffer - Update return format to include history dimension - Extend ring buffer size to accommodate history
- Fix reshape crash in read() for batched env subsets: use n_query_envs instead of global n_envs when reshaping history slices - Restore read_ground_truth() contract: return noise-free ground truth cache instead of delegating to read() which returns processed data - Fix _draw_debug() env index: use cache[env_idx] instead of hardcoded cache[0] for correct environment in multi-env scenes Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
- Adds entity.get_height_at(x, y) to query terrain height at world position using bilinear interpolation from height field - Adds entity.get_normal_at(x, y) to compute surface normal at world position from height field gradient - Both methods handle boundary conditions gracefully Closes Genesis-Embodied-AI#2094
…se transform - Fix coordinate indexing: hf[y,x] -> hf[x,y] since heightfield is stored as [row, col] where row corresponds to x - Add pose transformation: convert world coords to terrain local frame using inv_transform_by_trans_quat(terrain_pos, terrain_quat) - Transform normals back to world frame with transform_by_quat - Height now includes terrain z-offset (terrain_pos[2]) Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
…force sensor Terrain get_height_at / get_normal_at: - Remove inner edge-strip guard that returned non-zero height for positions just beyond the last grid node (x slightly > (N-1)*h_scale); now returns terrain_pos[2] consistently for all out-of-bounds positions - Switch from bilinear to triangle-based (barycentric) interpolation matching the physics mesh diagonal convention in convert_heightfield_to_watertight_trimesh: cell split along (x0,y0)-(x1,y1) diagonal, upper-left triangle when tx<=ty, lower-right when tx>ty; ensures queried heights/normals lie on simulated surface ContactForceSensor read() / read_ground_truth(): - Fix history tensor stacking: use dim=0 for non-batched (n_envs==0) to produce (history_length, 3) instead of (3, history_length); use dim=1 for batched to produce (n_envs, history_length, 3) — previous code transposed the result - Remove dead `if envs_idx is None` branch after _sanitize_envs_idx() which always returns a tensor - Clarify docstring: history buffer stores ground-truth values since the underlying ring buffer is written before noise/delay post-processing
…istory Terrain get_height_at / get_normal_at (test_rigid_physics.py): - test_terrain_get_height_flat: flat HF returns constant height+offset everywhere, normal is world-up - test_terrain_get_height_ramp_x: parametrized over 3 (shape, scale) combos; verifies linear height and correct slope direction for pure x-ramp - test_terrain_get_height_non_symmetric: non-square HF (10x20) with distinct slopes on each axis catches axis-swap bugs; includes analytic reference - test_terrain_get_height_triangle_consistency: at the diagonal (tx==ty==0.5) both triangle formulas must agree - test_terrain_get_height_out_of_bounds: 6 cases covering negative coords, positions just beyond last node (edge-strip), and far-outside; all must return pos_z, not edge height - test_terrain_get_height_pose_transform: parametrized over translated, yaw-45, and combined poses; world-space query of a known vertex must recover its height - test_terrain_get_normal_unit_length_and_direction: unit length + z>0 on random HF - test_terrain_triangle_selection_both_halves: direct validation of upper-left and lower-right triangle formulas using known corner heights; checks normals differ between triangles ContactForce history_length (test_sensors.py): - test_contact_force_history_length_1_shape: default history produces (3,) / (n,3) - test_contact_force_history_shape: history_length 2 and 5, n_envs 0 and 2; verifies both read() and read_ground_truth() shapes - test_contact_force_history_no_transpose: index [0] is most-recent, catches the dim=1 stacking bug (3, L) vs correct (L, 3) - test_contact_force_history_envs_idx_subset: requesting a 1-env subset of a batched scene must not raise and must return (1, L, 3) - test_contact_force_history_length_1_ground_truth_is_noiseless: with noise=5.0, read_ground_truth() must match physics; read() must differ - test_contact_force_history_before_and_after_contact: both history_length 1 and 3; GT force is zero before contact and non-zero after
rigid_entity.py — get_height_at / get_normal_at: - Add env_idx: int | None = None parameter to both methods; pass through to links[0].get_pos(env_idx).reshape((3,)) and get_quat(env_idx).reshape((4,)) so batched scenes can specify which env to query; reshape handles both the non-batched (3,) and batched (1,3) return shapes from get_pos/get_quat contact_force.py: - read_ground_truth(): move get_cloned_from_cache(is_ground_truth=True) clone inside the history_length==1 branch — was unconditionally allocating a GPU tensor even when history_length>1 where it was never used - Remove ContactForceSensorMetadata.history_length field: was updated in build() but read() / read_ground_truth() always use self._options.history_length directly, making the metadata field dead code - Remove corresponding self._shared_metadata.history_length update in build() tests/test_rigid_physics.py: - test_terrain_get_height_out_of_bounds: replace wrong edge-strip case (terrain_x_max - 0.01) which is actually IN bounds (x0=n_rows-2, x1=n_rows-1 are both valid) with (terrain_x_max, 0.5) which correctly triggers OOB (x0=n_rows-1, x1=n_rows >= hf.shape[0]); previous case would have failed - test_terrain_get_height_non_symmetric: remove dead variables h_swapped_x and h_swapped_y (same expression assigned twice, never read)
Dead import left after prior refactoring passes; cleaned up as part of ongoing PR review cycle.
moviepy >= 2.x drops the last frame when writing videos. Replace with PyAV which correctly flushes all frames via stream.encode(None). Fall back to moviepy when av is not installed. Fixes Genesis-Embodied-AI#1635
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Pull request overview
This PR updates Genesis’ video export utility to avoid the last-frame drop seen with moviepy ≥2.x by preferentially using PyAV in gs.tools.animate() (with a moviepy fallback), and also introduces new sensor/terrain functionality plus accompanying tests.
Changes:
- Replace the moviepy-based
gs.tools.animate()implementation with a PyAV-based writer that flushes frames viastream.encode(None); fall back to moviepy if PyAV isn’t installed. - Add
history_lengthsupport for theContactForcesensor, including SensorManager ring-buffer sizing changes. - Add
RigidEntity.get_height_at()/get_normal_at()terrain query APIs and extensive tests for terrain interpolation/normal behavior.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
genesis/utils/tools.py |
Switches animate() to PyAV encoding with moviepy fallback. |
genesis/options/sensors/options.py |
Adds history_length option to ContactForce sensor options. |
genesis/engine/sensors/sensor_manager.py |
Updates ring-buffer sizing to account for sensor history length. |
genesis/engine/sensors/contact_force.py |
Implements history reads for ContactForce and adjusts debug drawing to avoid history-shaped reads. |
genesis/engine/entities/rigid_entity/rigid_entity.py |
Adds terrain height/normal query helpers with pose handling and OOB behavior. |
tests/test_sensors.py |
Adds tests for ContactForce.history_length shape/order and cache-size behavior. |
tests/test_rigid_physics.py |
Adds tests for terrain get_height_at / get_normal_at correctness and edge cases. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
183
to
+188
| def __init__(self, options: ContactForceSensorOptions, sensor_idx: int, sensor_manager: "SensorManager"): | ||
| super().__init__(options, sensor_idx, sensor_manager) | ||
|
|
||
| self.debug_object: "Mesh" | None = None | ||
| self.debug_object: Mesh | None = None | ||
|
|
||
| def build(self): | ||
| super().build() | ||
| def _read_history(self, envs_idx: torch.Tensor) -> torch.Tensor: |
Comment on lines
+36
to
+61
| first = imgs[0] | ||
| if not isinstance(first, np.ndarray): | ||
| first = np.array(first) | ||
| height, width = first.shape[:2] | ||
| is_color = first.ndim == 3 and first.shape[2] == 3 | ||
|
|
||
| container = av.open(filename, mode="w") | ||
| stream = container.add_stream("libx264", rate=fps) | ||
| stream.width = width | ||
| stream.height = height | ||
| stream.pix_fmt = "yuv420p" | ||
| stream.codec_context.options = {"preset": "ultrafast"} | ||
|
|
||
| fmt = "rgb24" if is_color else "gray8" | ||
| video_frame = av.VideoFrame(width, height, fmt) | ||
| frame_plane = video_frame.planes[0] | ||
| if is_color: | ||
| buf = np.asarray(memoryview(frame_plane)).reshape((-1, frame_plane.line_size // 3, 3)) | ||
| else: | ||
| buf = np.asarray(memoryview(frame_plane)).reshape((-1, frame_plane.line_size)) | ||
|
|
||
| for img in imgs: | ||
| if not isinstance(img, np.ndarray): | ||
| img = np.array(img) | ||
| img = img.astype(np.uint8) | ||
| buf[: img.shape[0], : img.shape[1]] = img |
Comment on lines
+42
to
+68
| container = av.open(filename, mode="w") | ||
| stream = container.add_stream("libx264", rate=fps) | ||
| stream.width = width | ||
| stream.height = height | ||
| stream.pix_fmt = "yuv420p" | ||
| stream.codec_context.options = {"preset": "ultrafast"} | ||
|
|
||
| fmt = "rgb24" if is_color else "gray8" | ||
| video_frame = av.VideoFrame(width, height, fmt) | ||
| frame_plane = video_frame.planes[0] | ||
| if is_color: | ||
| buf = np.asarray(memoryview(frame_plane)).reshape((-1, frame_plane.line_size // 3, 3)) | ||
| else: | ||
| buf = np.asarray(memoryview(frame_plane)).reshape((-1, frame_plane.line_size)) | ||
|
|
||
| for img in imgs: | ||
| if not isinstance(img, np.ndarray): | ||
| img = np.array(img) | ||
| img = img.astype(np.uint8) | ||
| buf[: img.shape[0], : img.shape[1]] = img | ||
| for packet in stream.encode(video_frame): | ||
| container.mux(packet) | ||
|
|
||
| for packet in stream.encode(None): | ||
| container.mux(packet) | ||
| container.close() | ||
|
|
| max_buffer_len = max(max_buffer_len, sensor._delay_ts + 1) | ||
|
|
||
| history_length = getattr(sensor._options, "history_length", 1) | ||
| max_buffer_len = max(max_buffer_len, sensor._delay_ts + 1, history_length) |
Comment on lines
12
to
+35
| def animate(imgs, filename=None, fps=60): | ||
| """ | ||
| Create a video from a list of images. | ||
|
|
||
| Args: | ||
| imgs (list): List of input images. | ||
| filename (str, optional): Name of the output video file. If not provided, the name will be default to the name of the caller file, with a timestamp and '.mp4' extension. | ||
| """ | ||
| assert isinstance(imgs, list) | ||
| if len(imgs) == 0: | ||
| gs.logger.warning("No image to save.") | ||
| return | ||
|
|
||
| if filename is None: | ||
| caller_file = inspect.stack()[-1].filename | ||
| # caller file + timestamp + .mp4 | ||
| filename = os.path.splitext(os.path.basename(caller_file))[0] + f"_{time.strftime('%Y%m%d_%H%M%S')}.mp4" | ||
| os.makedirs(os.path.abspath(os.path.dirname(filename)), exist_ok=True) | ||
|
|
||
| gs.logger.info(f'Saving video to ~<"{filename}">~...') | ||
| from moviepy import ImageSequenceClip | ||
|
|
||
| imgs = ImageSequenceClip(imgs, fps=fps) | ||
| imgs.write_videofile( | ||
| filename, | ||
| fps=fps, | ||
| logger=None, | ||
| codec="libx264", | ||
| preset="ultrafast", | ||
| # ffmpeg_params=["-crf", "0"], | ||
| ) | ||
|
|
||
| try: | ||
| import av | ||
|
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0c0de887eb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| history_data = [] | ||
| for i in range(history_length): | ||
| hist = buffered_data.at(i, envs_idx, cache_slice) |
There was a problem hiding this comment.
When history_length > 1, both read() and read_ground_truth() call _read_history(), which always reads from self._manager._buffered_data. That buffer is only written in _update_shared_cache, but SensorManager skips _update_shared_cache when all sensors of this class set update_ground_truth_only=True, so history reads return stale/zero/uninitialized buffer contents instead of real contact-force history in that configuration.
Useful? React with 👍 / 👎.
| if not isinstance(first, np.ndarray): | ||
| first = np.array(first) | ||
| height, width = first.shape[:2] | ||
| is_color = first.ndim == 3 and first.shape[2] == 3 |
There was a problem hiding this comment.
The PyAV path only recognizes color when shape[2] == 3; RGBA frames (shape[2] == 4) are misclassified as grayscale, so fmt becomes gray8 and buf is 2D while frame assignment later uses a 3D array, causing a broadcast error. This breaks gs.tools.animate() for common PNG-with-alpha inputs (including those loaded via PIL.Image.open).
Useful? React with 👍 / 👎.
3 tasks
Author
|
Closing in favour of #2705, which is a clean replacement rebased directly on main. This PR inadvertently included unrelated terrain query and ContactForce sensor changes in its diff because the branch was created from rather than . The new PR also addresses the Copilot/Codex review comments: RGBA frame handling and the missing try/finally for container cleanup. |
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.
Summary
moviepy-based implementation ings.tools.animate()with PyAV (av)stream.encode(None)before closing the containeravis not installedProblem
In moviepy >= 2.x,
ImageSequenceClip.write_videofile()drops the last frame of the video. This was reproducible and confirmed by the community in #1635. The maintainer noted that "A new recorder based on PyAV has been implemented, but camera needs to be updated to leverage this new capability" — this PR wires up that fix foranimate()/camera.stop_recording().Test plan
start_recording()/stop_recording()avis not installed (falls back to moviepy without error)Closes #1635