From ef7f7b4f7df8c073b8991050a3702ebb01bd7eba Mon Sep 17 00:00:00 2001 From: BartR Date: Sat, 27 Dec 2025 17:32:20 +0100 Subject: [PATCH 1/3] wip: videos on iteration trigger --- genesis_forge/wrappers/__init__.py | 3 +- genesis_forge/wrappers/video.py | 83 ++++++++++++++++++++++++++---- 2 files changed, 76 insertions(+), 10 deletions(-) diff --git a/genesis_forge/wrappers/__init__.py b/genesis_forge/wrappers/__init__.py index 30d1f54..0365bf2 100644 --- a/genesis_forge/wrappers/__init__.py +++ b/genesis_forge/wrappers/__init__.py @@ -1,9 +1,10 @@ -from .video import VideoWrapper +from .video import VideoWrapper, capped_cubic_iteration_trigger from .rsl_rl import RslRlWrapper from .skrl import SkrlEnvWapper __all__ = [ "VideoWrapper", + "capped_cubic_iteration_trigger", "RslRlWrapper", "SkrlEnvWapper", ] diff --git a/genesis_forge/wrappers/video.py b/genesis_forge/wrappers/video.py index b5deb11..0869505 100644 --- a/genesis_forge/wrappers/video.py +++ b/genesis_forge/wrappers/video.py @@ -5,7 +5,7 @@ import torch from genesis_forge.genesis_env import GenesisEnv from genesis_forge.wrappers.wrapper import Wrapper -from typing import Tuple, Any, Callable, Literal, TYPE_CHECKING +from typing import Tuple, Any, Callable, Literal, Optional, TYPE_CHECKING if TYPE_CHECKING: from genesis.vis.camera import Camera @@ -30,17 +30,34 @@ def capped_cubic_episode_trigger(episode_id: int) -> bool: return episode_id % 1000 == 0 +def capped_cubic_iteration_trigger(iteration: int) -> bool: + """The default iteration trigger. + + This function will trigger recordings at the iteration indices 0, 1, 8, 27, ..., :math:`k^3`, ..., 729, 1000, 2000, 3000, ... + + Args: + iteration: The iteration number + + Returns: + If to apply a video schedule number + """ + if iteration < 1000: + return int(round(iteration ** (1.0 / 3))) ** 3 == iteration + else: + return iteration % 1000 == 0 + + class VideoWrapper(Wrapper): """ - Automatically record videos during training at a regular step or episode intervals. + Automatically record videos during training at a regular step, episode, or iteration intervals. Based on the RecordVideo wrapper from Gymnasium: https://gymnasium.farama.org/main/api/wrappers/misc_wrappers/#gymnasium.wrappers.RecordVideo Recordings will be made from a dedicated camera, which you need to add to your environment (see the example below). - To control how frequently recordings are made specify **either** ``episode_trigger`` **or** ``step_trigger`` (not both). + To control how frequently recordings are made specify **one of** ``episode_trigger``, ``step_trigger``, or ``iteration_trigger``. They should be functions returning a boolean that indicates whether a recording should be started at the - current episode or step, respectively. If neither :attr:`episode_trigger` nor ``step_trigger`` is passed, + current episode, step, or iteration, respectively. If none are passed, a default ``episode_trigger`` will be used, which records at the episode indices 0, 1, 8, 27, ..., :math:`k^3`, ..., 729, 1000, 2000, 3000,. Args: @@ -48,12 +65,16 @@ class VideoWrapper(Wrapper): camera_attr: The attribute of the base environment that contains the camera to use for recording. episode_trigger: Function that accepts an episode count integer and returns ``True`` if a recording should be started at this episode step_trigger: Function that accepts a step count integer and returns ``True`` if a recording should be started at this step + iteration_trigger: Function that accepts an iteration count integer and returns ``True`` if a recording should be started at this iteration. + Requires ``num_steps_per_env`` to be set to calculate iteration from step count. + num_steps_per_env: Number of steps per environment per iteration. Required when using ``iteration_trigger``. + initial_iteration: Initial iteration offset (useful when resuming training). video_length_sec: Length of each video, in seconds. out_dir: Directory to save the videos to. fps: Frames per second for the video. env_idx: If triggering on episode, this is the index of the environment to be counting episodes for. filename: The filename for the video. - If None, the video will automatically be named for the current step. + If None, the video will automatically be named for the current step (or iteration if using iteration_trigger). If defined, each video will overwrite the previous video with this name. Example:: @@ -88,6 +109,17 @@ def train(): out_dir="./videos", step_trigger=lambda step: step % 1500 == 0 ) + + Record every 50 iterations (useful for RL training):: + + env = MyEnv() + env = VideoWrapper( + env, + camera_attr="camera", + out_dir="./videos", + iteration_trigger=lambda it: it % 50 == 0, + num_steps_per_env=24, # Should match your training config + ) """ def __init__( @@ -97,6 +129,9 @@ def __init__( video_length_sec: int = 8, episode_trigger: Callable[[int], bool] | None = None, step_trigger: Callable[[int], bool] | None = None, + iteration_trigger: Callable[[int], bool] | None = None, + num_steps_per_env: Optional[int] = None, + initial_iteration: int = 0, out_dir: str = "./videos", fps: int = 60, env_idx: int = 0, @@ -109,7 +144,10 @@ def __init__( self._logging: bool = logging self._current_step: int = 0 self._current_episode: int = 0 + self._current_iteration: int = initial_iteration + self._last_triggered_iteration: int = initial_iteration - 1 self._recording_start_step: int = 0 + self._recording_start_iteration: int = 0 self._recording_stop_step: int = 0 self._record_final_episode = record_final_episode self._has_recording_buffer = False @@ -128,14 +166,23 @@ def __init__( self._actual_fps = round(1.0 / self.dt / self._steps_per_frame) self._env_idx = env_idx - if episode_trigger is None and step_trigger is None: + # Iteration tracking + self._num_steps_per_env = num_steps_per_env + self._initial_iteration = initial_iteration + + # Validate trigger configuration + if iteration_trigger is not None and num_steps_per_env is None: + raise ValueError("num_steps_per_env must be specified when using iteration_trigger") + + if episode_trigger is None and step_trigger is None and iteration_trigger is None: episode_trigger = capped_cubic_episode_trigger - trigger_count = sum(x is not None for x in [episode_trigger, step_trigger]) - assert trigger_count == 1, "Must specify only one trigger" + trigger_count = sum(x is not None for x in [episode_trigger, step_trigger, iteration_trigger]) + assert trigger_count == 1, "Must specify only one trigger (episode_trigger, step_trigger, or iteration_trigger)" self.episode_trigger = episode_trigger self.step_trigger = step_trigger + self.iteration_trigger = iteration_trigger os.makedirs(self._out_dir, exist_ok=True) @@ -166,6 +213,10 @@ def step( extras, ) = super().step(actions) + # Update iteration count if using iteration trigger + if self._num_steps_per_env is not None: + self._current_iteration = self._initial_iteration + (self._current_step // self._num_steps_per_env) + self._check_recording_trigger() if self._current_step % self._steps_per_frame == 0: self._cam.render() @@ -208,6 +259,7 @@ def start_recording(self, type: RecordingType = "active"): self._has_recording_buffer = False self._recording_type = type self._recording_start_step = self._current_step + self._recording_start_iteration = self._current_iteration self._recording_stop_step = self._current_step + self._video_length_steps self._cam.start_recording() @@ -220,7 +272,13 @@ def finish_recording(self): # Save recording if self._recording_type == "active": - filename = self._filename or f"{self._recording_start_step}.mp4" + # Use iteration-based filename if iteration_trigger is active + if self._filename: + filename = self._filename + elif self.iteration_trigger is not None: + filename = f"iter_{self._recording_start_iteration}.mp4" + else: + filename = f"{self._recording_start_step}.mp4" filepath = os.path.join(self._out_dir, filename) if self._logging: print(f"Saving recording to {filepath}") @@ -237,12 +295,19 @@ def finish_recording(self): def _check_recording_trigger(self) -> bool: """Check if a recording should be started""" + record = False if self._is_recording and self._recording_type == "active": record = False elif self.episode_trigger is not None: record = self.episode_trigger(self._current_episode) elif self.step_trigger is not None: record = self.step_trigger(self._current_step) + elif self.iteration_trigger is not None: + # Only trigger once per iteration (at the start of each new iteration) + if self._current_iteration > self._last_triggered_iteration: + if self.iteration_trigger(self._current_iteration): + record = True + self._last_triggered_iteration = self._current_iteration if record: self.start_recording() From 6166db4389d6a10c99c0f60c3e34d284e42bbb2d Mon Sep 17 00:00:00 2001 From: BartR Date: Tue, 30 Dec 2025 12:47:36 +0100 Subject: [PATCH 2/3] add new gamepad model support --- genesis_forge/gamepads/__init__.py | 3 +- genesis_forge/gamepads/dragonrise.py | 33 +++++++++++++++++++ genesis_forge/gamepads/gamepad.py | 5 ++- .../managers/command/command_manager.py | 6 ++++ 4 files changed, 45 insertions(+), 2 deletions(-) create mode 100644 genesis_forge/gamepads/dragonrise.py diff --git a/genesis_forge/gamepads/__init__.py b/genesis_forge/gamepads/__init__.py index 1e06f93..a01089b 100644 --- a/genesis_forge/gamepads/__init__.py +++ b/genesis_forge/gamepads/__init__.py @@ -1,5 +1,6 @@ from .gamepad import Gamepad from .logitech import LOGITECH_F710_CONFIG, LOGITECH_F310_CONFIG +from .dragonrise import DRAGONRISE_CONFIG from .config import GamepadState -__all__ = ["Gamepad", "GamepadState", "LOGITECH_F710_CONFIG", "LOGITECH_F310_CONFIG"] +__all__ = ["Gamepad", "GamepadState", "LOGITECH_F710_CONFIG", "LOGITECH_F310_CONFIG", "DRAGONRISE_CONFIG"] diff --git a/genesis_forge/gamepads/dragonrise.py b/genesis_forge/gamepads/dragonrise.py new file mode 100644 index 0000000..41fbb16 --- /dev/null +++ b/genesis_forge/gamepads/dragonrise.py @@ -0,0 +1,33 @@ +"""DragonRise Inc. Generic USB Joystick configuration. + +Common cheap USB gamepad found on Amazon/AliExpress. +Vendor ID: 0x0079, Product ID: 0x0006 +""" + +from .config import GamepadConfig + +VENDOR_ID = 0x0079 +PRODUCT_ID = 0x0006 + +# HID data mapping (determined via testing): +# Byte 0: Left stick X (0=left, 128=center, 255=right) +# Byte 1: Left stick Y (0=up/forward, ~109=center, 255=down/back) +# Byte 2: Right stick X (0=left, ~125=center, 255=right) +# Byte 3: Right stick Y (128=center) +# Byte 4: Triggers/Z-axis +# Byte 5: D-pad and face buttons (lower nibble = d-pad, upper nibble = buttons) +# Byte 6: Shoulder buttons +# Byte 7: Mode/special buttons + +DRAGONRISE_CONFIG: GamepadConfig = { + "name": "DragonRise Generic USB Joystick", + "vendor_id": VENDOR_ID, + "product_id": PRODUCT_ID, + "mapping": [ + # Analog sticks + {"axis": 0, "data": 0}, # Left stick X + {"axis": 1, "data": 1}, # Left stick Y + {"axis": 2, "data": 2}, # Right stick X + {"axis": 3, "data": 3}, # Right stick Y + ], +} diff --git a/genesis_forge/gamepads/gamepad.py b/genesis_forge/gamepads/gamepad.py index ec2999b..b5bf2f7 100644 --- a/genesis_forge/gamepads/gamepad.py +++ b/genesis_forge/gamepads/gamepad.py @@ -5,10 +5,12 @@ from .config import GamepadConfig, GamepadState from .logitech import LOGITECH_F710_CONFIG, LOGITECH_F310_CONFIG +from .dragonrise import DRAGONRISE_CONFIG GAMEPAD_CONFIGS = [ LOGITECH_F710_CONFIG, LOGITECH_F310_CONFIG, + DRAGONRISE_CONFIG, ] @@ -131,7 +133,8 @@ def _read_loop(self): """ while self.is_running: try: - data = self._device.read(64) + # Use timeout (10ms) to prevent blocking - important for macOS + data = self._device.read(64, timeout_ms=10) if data: try: self._state = self._parse_data(data) diff --git a/genesis_forge/managers/command/command_manager.py b/genesis_forge/managers/command/command_manager.py index 433c517..3ce27d4 100644 --- a/genesis_forge/managers/command/command_manager.py +++ b/genesis_forge/managers/command/command_manager.py @@ -381,10 +381,16 @@ def resample_command(self, env_ids: list[int]): def _gamepad_axis_command(self, step_count: int) -> torch.Tensor: """ Get the command from the gamepad. + Uses per-step caching to avoid redundant reads (command may be accessed multiple times per step). """ if self._gamepad_cfg is None: return self._gamepad_axis_command_buffer + # Cache: only read gamepad once per step + if hasattr(self, '_last_gamepad_step') and self._last_gamepad_step == step_count: + return self._gamepad_axis_command_buffer + self._last_gamepad_step = step_count + gamepad = self._gamepad_cfg["gamepad"] axis_map = self._gamepad_cfg["axis_map"] From 77c905d06db615b9da7efc3e47d2401a88d9b560 Mon Sep 17 00:00:00 2001 From: BartR Date: Tue, 30 Dec 2025 12:48:07 +0100 Subject: [PATCH 3/3] video recording mods --- genesis_forge/wrappers/video.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/genesis_forge/wrappers/video.py b/genesis_forge/wrappers/video.py index 0869505..a8ac126 100644 --- a/genesis_forge/wrappers/video.py +++ b/genesis_forge/wrappers/video.py @@ -63,13 +63,15 @@ class VideoWrapper(Wrapper): Args: env: GenesisEnv camera_attr: The attribute of the base environment that contains the camera to use for recording. + video_length_sec: Length of each video, in seconds. Use this OR ``video_length_iterations``. + video_length_iterations: Length of each video, in iterations. Useful for iteration_trigger to ensure + recording duration doesn't exceed trigger interval. Requires ``num_steps_per_env``. episode_trigger: Function that accepts an episode count integer and returns ``True`` if a recording should be started at this episode step_trigger: Function that accepts a step count integer and returns ``True`` if a recording should be started at this step iteration_trigger: Function that accepts an iteration count integer and returns ``True`` if a recording should be started at this iteration. Requires ``num_steps_per_env`` to be set to calculate iteration from step count. - num_steps_per_env: Number of steps per environment per iteration. Required when using ``iteration_trigger``. + num_steps_per_env: Number of steps per environment per iteration. Required when using ``iteration_trigger`` or ``video_length_iterations``. initial_iteration: Initial iteration offset (useful when resuming training). - video_length_sec: Length of each video, in seconds. out_dir: Directory to save the videos to. fps: Frames per second for the video. env_idx: If triggering on episode, this is the index of the environment to be counting episodes for. @@ -119,6 +121,7 @@ def train(): out_dir="./videos", iteration_trigger=lambda it: it % 50 == 0, num_steps_per_env=24, # Should match your training config + video_length_iterations=1, # Record 1 iteration per video ) """ @@ -126,7 +129,8 @@ def __init__( self, env: GenesisEnv, camera_attr: str = "camera", - video_length_sec: int = 8, + video_length_sec: Optional[int] = None, + video_length_iterations: Optional[int] = None, episode_trigger: Callable[[int], bool] | None = None, step_trigger: Callable[[int], bool] | None = None, iteration_trigger: Callable[[int], bool] | None = None, @@ -161,7 +165,6 @@ def __init__( self._camera_attr = camera_attr self._out_dir = out_dir self._filename = filename - self._video_length_steps = math.ceil(video_length_sec / self.dt) self._steps_per_frame = round(1.0 / fps / self.dt) self._actual_fps = round(1.0 / self.dt / self._steps_per_frame) self._env_idx = env_idx @@ -169,6 +172,18 @@ def __init__( # Iteration tracking self._num_steps_per_env = num_steps_per_env self._initial_iteration = initial_iteration + self._video_length_iterations = video_length_iterations + + # Calculate video length in steps + if video_length_iterations is not None: + if num_steps_per_env is None: + raise ValueError("num_steps_per_env must be specified when using video_length_iterations") + self._video_length_steps = video_length_iterations * num_steps_per_env + elif video_length_sec is not None: + self._video_length_steps = math.ceil(video_length_sec / self.dt) + else: + # Default: 8 seconds + self._video_length_steps = math.ceil(8 / self.dt) # Validate trigger configuration if iteration_trigger is not None and num_steps_per_env is None: