Skip to content

Fix(record): save camera recording on episode reset #208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuecideng merged 1 commit into from
Mar 31, 2026
+29 −9
Merged

Fix(record): save camera recording on episode reset #208

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions embodichain/lab/gym/envs/embodied_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -515,6 +515,15 @@ def _initialize_episode(
env_ids=successful_env_ids.nonzero(as_tuple=True)[0],
)

# Save recorded camera data before resetting
if self.cfg.events and self.event_manager is not None:
from embodichain.lab.gym.envs.managers.record import record_camera_data

for mode_cfgs in self.event_manager._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if isinstance(functor_cfg.func, record_camera_data):
functor_cfg.func.save_and_clear()
Comment on lines +518 to +525
Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change introduces new episode-finalization behavior (calling save_and_clear() during resets), but there are no tests asserting that recordings are saved/cleared at the right time. Given BaseEnv.step() resets only a subset of envs on dones, adding a unit test around _initialize_episode(reset_ids=...) would help prevent regressions (e.g., ensure save/clear is only triggered on full resets or the intended condition).

Copilot uses AI. Check for mistakes.
Comment on lines +518 to +525
Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_initialize_episode() is called for partial resets as well (see BaseEnv.step() auto-resets only dones envs). Calling save_and_clear() here unconditionally will clear the shared frame buffer even when only a subset of envs is being reset, which can truncate videos and drop frames for envs that are still mid-episode. Consider only calling save_and_clear() when resetting all envs (e.g., env_ids is None or len(env_ids)==self.num_envs), or rework the recorder to buffer/save per-env when partial resets occur.

Suggested change
# Save recorded camera data before resetting
if self.cfg.events and self.event_manager is not None:
from embodichain.lab.gym.envs.managers.record import record_camera_data
for mode_cfgs in self.event_manager._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if isinstance(functor_cfg.func, record_camera_data):
functor_cfg.func.save_and_clear()
# Save recorded camera data before resetting.
# Only clear the shared recorder buffer when all environments are being reset.
if self.cfg.events and self.event_manager is not None:
from embodichain.lab.gym.envs.managers.record import record_camera_data
# A global reset occurs when env_ids is None, or when the number of
# environments to process matches the total number of environments.
if env_ids is None or len(env_ids_to_process) == self.num_envs:
for mode_cfgs in self.event_manager._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if isinstance(functor_cfg.func, record_camera_data):
functor_cfg.func.save_and_clear()

Copilot uses AI. Check for mistakes.

# Clear episode buffers and reset success status for environments being reset
if self.rollout_buffer is not None and self._rollout_buffer_mode != "rl":
self.current_rollout_step = 0
Expand Down
29 changes: 20 additions & 9 deletions embodichain/lab/gym/envs/managers/record.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
if getattr(env.sim, "is_rt_enabled", False):
env.add_camera_group_id(self.camera.group_id)

self._save_path = cfg.params.get("save_path", "./outputs/videos")
self._current_episode = 0
self._frames: List[np.ndarray] = []

Expand Down Expand Up @@ -124,6 +125,21 @@ def _draw_frames_into_one_image(self, frames: torch.Tensor) -> torch.Tensor:

return result

def save_and_clear(self) -> None:
"""Save recorded frames as video and clear the buffer.

This method is called from :meth:`EmbodiedEnv._initialize_episode` to ensure
frames are saved before the episode is reset. This avoids the issue where the
final episode's frames are lost because the save previously relied on detecting
a reset inside :meth:`__call__`.
"""
if len(self._frames) > 0:
video_name = f"episode_{self._current_episode}_{self._name}"
images_to_video(self._frames, self._save_path, video_name, fps=20)

self._current_episode += 1
self._frames = []

def __call__(
self,
env: EmbodiedEnv,
Expand All @@ -142,15 +158,6 @@ def __call__(
max_env_num: int = 16,
save_path: str = "./outputs/videos",
):
Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

save_path is still part of the functor call signature, but it is no longer used (saving now uses self._save_path captured in __init__). This makes runtime overrides of save_path ineffective and is inconsistent with record_camera_data_async (which still uses the passed save_path). Either remove save_path from the signature/params (and update configs accordingly) or set self._save_path = save_path inside __call__ so both code paths respect the same input.

Suggested change
):
):
# Ensure per-call save_path overrides the instance-level setting used by save_and_clear
self._save_path = save_path

Copilot uses AI. Check for mistakes.
# TODO: the current implementation will lost the final episode frames recording.
# Check if the frames should be saved for the current episode
if env.elapsed_steps.sum().item() == 0 and len(self._frames) > 0:
video_name = f"episode_{self._current_episode}_{self._name}"
images_to_video(self._frames, save_path, video_name, fps=20)

self._current_episode += 1
self._frames = []

self.camera.update(fetch_only=self.camera.is_rt_enabled)
data = self.camera.get_data()
rgb = data["color"]
Comment on lines 158 to 163
Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

max_env_num is intended to limit how many env frames are used, but the current tiling logic later computes num_frames using max(rgb.shape[0], max_env_num), which won’t cap anything when rgb.shape[0] > max_env_num. Consider switching that computation to min(...) so max_env_num reliably bounds the grid size and avoids unnecessary work.

Copilot uses AI. Check for mistakes.
Expand All @@ -170,6 +177,10 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
self._frames_list = [[] for _ in range(self._num_envs)]
self._ep_idx = [0 for _ in range(self._num_envs)]

def save_and_clear(self) -> None:
"""No-op for the async variant; saving is handled inside :meth:`__call__`."""
pass

def __call__(
self,
env: EmbodiedEnv,
Expand Down
Loading