Add support for writing joystick data as action#13
Conversation
There was a problem hiding this comment.
Pull request overview
Adds joystick stream recording to the dataset writer so joystick inputs can be persisted alongside other action signals (GH-7).
Changes:
- Extend
Episodeto collect joystick x/y/button values and timestamps. - Persist joystick streams to parquet under
episodes/<id>/action/joystick/. - Ingest
joystick_*events in the main event loop and append them into the active episode buffers.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def _write_joystick_buttons(self, output_path, timestamps, buttons): | ||
| output_path.parent.mkdir(parents=True, exist_ok=True) | ||
| table = pa.table( | ||
| { | ||
| "timestamp": pa.array(timestamps, type=pa.timestamp("ns")), | ||
| "value": pa.concat_arrays(buttons), | ||
| } | ||
| ) | ||
| pq.write_table(table, output_path) | ||
|
|
||
|
|
There was a problem hiding this comment.
_write_joystick_buttons is currently unused (buttons are written via _write_joystick_values). This adds dead/duplicated code; either remove _write_joystick_buttons or use it consistently for the button stream to avoid confusion and keep the writer surface minimal.
| def _write_joystick_buttons(self, output_path, timestamps, buttons): | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| table = pa.table( | |
| { | |
| "timestamp": pa.array(timestamps, type=pa.timestamp("ns")), | |
| "value": pa.concat_arrays(buttons), | |
| } | |
| ) | |
| pq.write_table(table, output_path) |
| elif event_id.startswith("joystick_"): | ||
| # joystick_x -> | ||
| # joystick_xs | ||
| values_key = f"{event_id}s" | ||
| getattr(episode, values_key).append(event["value"]) | ||
| # joystick_x -> | ||
| # joystick_x_timestamps | ||
| timestamps_key = f"{event_id}_timestamps" | ||
| getattr(episode, timestamps_key).append(timestamp) |
There was a problem hiding this comment.
Now that joystick streams are being recorded under action/joystick/*, the generated metadata['frequencies'] structure never includes joystick inputs (only arm_ and camera_). Consider extending _collect_dynamic_metadata to populate an action.joystick section (and include joystick inputs in the node config) so dataset consumers can discover joystick action rates consistently.
| if self._episode.joystick_xs: | ||
| self._write_joystick_values( | ||
| self._base_directory / "action" / "joystick" / "x.parquet", | ||
| self._episode.joystick_x_timestamps, | ||
| self._episode.joystick_xs, | ||
| ) | ||
| if self._episode.joystick_ys: | ||
| self._write_joystick_values( | ||
| self._base_directory / "action" / "joystick" / "y.parquet", | ||
| self._episode.joystick_y_timestamps, | ||
| self._episode.joystick_ys, | ||
| ) | ||
| if self._episode.joystick_buttons: | ||
| self._write_joystick_values( | ||
| self._base_directory / "action" / "joystick" / "button.parquet", | ||
| self._episode.joystick_button_timestamps, | ||
| self._episode.joystick_buttons, | ||
| ) |
There was a problem hiding this comment.
Joystick action persistence is new behavior but there are no tests exercising EpisodeWriter.finish() for the joystick streams (parquet written, row counts matching timestamps, expected schema/types). Since this module already has pytest coverage for FrequencyDetector, adding a small tempfile-based test would help prevent regressions.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 1 out of 1 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| metadata["frequencies"][type]["arms"][side] = frequency | ||
| elif name.startswith("joystick_"): | ||
| target = name.split("_", 1)[1] | ||
| metadata["frequencies"]["joystick"][target] = frequency |
There was a problem hiding this comment.
In _collect_dynamic_metadata, joystick frequencies are initialized under metadata["frequencies"]["action"]["joystick"], but the assignment writes to metadata["frequencies"]["joystick"][target]. This will raise a KeyError once a joystick_* input is present, and it also places the data in the wrong section of the metadata. Write to metadata["frequencies"]["action"]["joystick"][target] instead (consistent with the initialized structure).
| metadata["frequencies"]["joystick"][target] = frequency | |
| metadata["frequencies"]["action"]["joystick"][target] = frequency |
| metadata["frequencies"][type]["arms"][side] = frequency | ||
| elif name.startswith("joystick_"): | ||
| target = name.split("_", 1)[1] | ||
| metadata["frequencies"]["joystick"][target] = frequency |
There was a problem hiding this comment.
The new joystick frequency mapping in _collect_dynamic_metadata isn’t covered by tests. Since this file already has pytest coverage for FrequencyDetector, consider adding a small test that builds a minimal node_config()/dataflow_descriptor() stub with a joystick_x (or similar) input and asserts the resulting metadata["frequencies"]["action"]["joystick"] entries, to prevent regressions like wrong nesting/KeyErrors.
| metadata["frequencies"]["joystick"][target] = frequency | |
| metadata["frequencies"]["action"]["joystick"][target] = frequency |
GitHub: fix GH-7