Skip to content

Add support for writing joystick data as action #13

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

Draft
kou wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add support for writing joystick data as action #13

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
47 changes: 47 additions & 0 deletions src/dora_openarm_dataset_recorder/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,12 @@ class Episode:
left_actions: ArrayLike = field(default_factory=list)
left_observation_timestamps: ArrayLike = field(default_factory=list)
left_observations: ArrayLike = field(default_factory=list)
joystick_x_timestamps: ArrayLike = field(default_factory=list)
joystick_xs: ArrayLike = field(default_factory=list)
joystick_y_timestamps: ArrayLike = field(default_factory=list)
joystick_ys: ArrayLike = field(default_factory=list)
joystick_button_timestamps: ArrayLike = field(default_factory=list)
joystick_buttons: ArrayLike = field(default_factory=list)


class EpisodeWriter:
Expand Down Expand Up @@ -94,6 +100,24 @@ def finish(self):
self._episode.left_observation_timestamps,
self._episode.left_observations,
)
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,
)
Comment on lines +103 to +120
Copy link

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

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

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.

Copilot uses AI. Check for mistakes.

def cancel(self):
"""Cancel this episode."""
Expand All @@ -110,6 +134,16 @@ def _write_positions(self, output_path, timestamps, positions):
)
pq.write_table(table, output_path)

def _write_joystick_values(self, output_path, timestamps, values):
output_path.parent.mkdir(parents=True, exist_ok=True)
table = pa.table(
{
"timestamp": pa.array(timestamps, type=pa.timestamp("ns")),
"value": pa.concat_arrays(values),
}
)
pq.write_table(table, output_path)


class DatasetWriter:
"""Write a dataset."""
Expand Down Expand Up @@ -196,6 +230,7 @@ def _collect_dynamic_metadata(metadata, args, node):
metadata["frequencies"] = {
"action": {
"arms": {},
"joystick": {},
},
"obs": {
"arms": {},
Expand All @@ -213,6 +248,9 @@ def _collect_dynamic_metadata(metadata, args, node):
if type == "observation":
type = "obs"
metadata["frequencies"][type]["arms"][side] = frequency
elif name.startswith("joystick_"):
target = name.split("_", 1)[1]
metadata["frequencies"]["joystick"][target] = frequency
Copy link

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

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

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).

Suggested change
metadata["frequencies"]["joystick"][target] = frequency
metadata["frequencies"]["action"]["joystick"][target] = frequency

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

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

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.

Suggested change
metadata["frequencies"]["joystick"][target] = frequency
metadata["frequencies"]["action"]["joystick"][target] = frequency

Copilot uses AI. Check for mistakes.
elif name.startswith("camera_"):
# camera_wrist_right -> wrist_right
camera_name = name.removeprefix("camera_")
Expand Down Expand Up @@ -329,6 +367,15 @@ def main():
image = event["value"]
format = event["metadata"]["encoding"]
episode_writer.write_camera_image(name, image, timestamp, format)
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)
Comment on lines +370 to +378
Copy link

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

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

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.

Copilot uses AI. Check for mistakes.


if __name__ == "__main__":
Expand Down
Loading