fix(imitation): improve collection safety and episode handling

[ruthwikdasyam]

Contribution path

• Small, safe change that does not need a tracking issue

Problem

Imitation collection performed unnecessary pose/TF lookup for semantically poseless streams, could stop its recorder before final producer events, and did not clearly report an episode left open at the end of a recording.

Quest disconnects could also leave stale controller engagement active, and invalid episode button mappings were silently ignored.

Solution

• Keep collection streams explicitly poseless and disable TF recording
• Stop camera, teleop, coordinator, and episode producers before the collection recorder
• Persist an explicit discard when shutdown interrupts an active episode
• Clear Quest controller state on the final client disconnect, publish a zero-button safe command, require fresh reconnect state, and reject unknown controller identities
• Validate episode button mappings at startup
• Extend dataprep inspection to report incomplete source episodes separately from saved/discarded episodes.
• Add a real SQLite-to-HDF5/LeRobot integration test

How to Test

Run the recording inspector against a collection database:

dimos dataprep inspect /path/to/session.db

Automated verification:

uv run pytest dimos/memory dimos/imitation dimos/teleop/quest -q

uv run mypy dimos/

uv run pre-commit run --all-files

Results: 541 tests passed, mypy passed across 999 source files, and all pre-commit hooks passed.

AI assistance

OpenCode with GPT-5.6 Sol was used extensively for codebase analysis, implementation, merge-conflict resolution, tests, and verification. The changes and generated tests were reviewed through focused diffs, type checking, pre-commit, and the affected test suites.

Checklist

• I have read and approved the CLA.

[codecov]

Codecov Report

❌ Patch coverage is 95.03546% with 28 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
dimos/teleop/quest/quest_teleop_module.py	84.84%	5 Missing and 5 partials ⚠️
dimos/imitation/collection/episode_monitor.py	88.57%	3 Missing and 5 partials ⚠️
dimos/imitation/test_datacollection_e2e.py	96.42%	5 Missing ⚠️
dimos/imitation/collection/recorder.py	66.66%	1 Missing and 1 partial ⚠️
dimos/teleop/quest/quest_extensions.py	83.33%	1 Missing and 1 partial ⚠️
dimos/imitation/dataprep/cli.py	0.00%	1 Missing ⚠️

@@            Coverage Diff             @@
##             main    #3497      +/-   ##
==========================================
+ Coverage   76.99%   77.18%   +0.18%     
==========================================
  Files        1260     1262       +2     
  Lines      119791   120299     +508     
  Branches    10603    10731     +128     
==========================================
+ Hits        92237    92852     +615     
+ Misses      24500    24368     -132     
- Partials     3054     3079      +25     

Flag	Coverage Δ
OS-ubuntu-24.04-arm	71.99% <77.12%> (+0.04%)	⬆️
OS-ubuntu-latest	73.89% <95.03%> (+0.20%)	⬆️
Py-3.10	73.88% <95.03%> (+0.20%)	⬆️
Py-3.11	73.88% <95.03%> (+0.20%)	⬆️
Py-3.12	73.88% <95.03%> (+0.20%)	⬆️
Py-3.13	73.88% <95.03%> (+0.19%)	⬆️
Py-3.14	73.88% <95.03%> (+0.19%)	⬆️
Py-3.14t	73.88% <95.03%> (+0.19%)	⬆️
SelfHosted-Large	30.03% <19.14%> (-0.06%)	⬇️
SelfHosted-Linux	35.42% <19.14%> (-0.08%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
dimos/cli/commands/dataprep.py	83.33% <ø> (ø)
dimos/imitation/collection/blueprint.py	88.23% <ø> (ø)
dimos/imitation/collection/test_blueprint.py	100.00% <100.00%> (ø)
dimos/imitation/collection/test_episode_monitor.py	99.35% <100.00%> (+0.48%)	⬆️
dimos/imitation/collection/test_recorder.py	100.00% <100.00%> (ø)
dimos/imitation/dataprep/build.py	91.56% <100.00%> (+62.15%)	⬆️
dimos/imitation/dataprep/core.py	92.42% <100.00%> (+5.12%)	⬆️
dimos/imitation/dataprep/test_core.py	99.44% <100.00%> (+0.06%)	⬆️
dimos/teleop/quest/test_quest_teleop_module.py	100.00% <100.00%> (ø)
dimos/imitation/dataprep/cli.py	19.04% <0.00%> (ø)
... and 5 more

... and 7 files with indirect coverage changes

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[greptile-apps]

Greptile Summary

The collection and dataprep updates improve interrupted-episode handling and recording inspection, while the Quest teleop update is intended to stop the Go2 on invalid controller input.

T-Rex validation blocked

The Quest controller safety check could not run because the local validation environment is missing the sse_starlette package required by a transitive web-server import. The runtime check stopped before it could submit controller input or observe the velocity command.

Confidence Score: 4/5

Do not merge until the invalid-controller stop behavior has been exercised in the runtime environment.

An invalid controller identity can leave a prior Go2 drive command active unless the safe-stop path is confirmed to run. The attempted runtime check was blocked before that path could execute.

Files Needing Attention: dimos/teleop/quest/quest_extensions.py

T-Rex Logs

What T-Rex did

• T-Rex attempted to run the quest-unknown-controller-validation.py script from /home/user/repo, preparing a valid left-controller Joy frame with a nonzero command, then an unknown-controller-id frame, and aiming to verify that ValueError propagates after the latest command is Twist.zero().
• The run was blocked during import by ModuleNotFoundError: No module named 'sse_starlette' from a transitive web-server import, preventing module construction and Joy-path execution, so no runtime result could be observed; the source also notes an exception handler that would publish a safe zero before re-raising, but no runtime result was available.

_{Ran code and verified through T-Rex}

_{Reviews (5): Last reviewed commit: "fix(teleop): stop Go2 on unknown control..." | Re-trigger Greptile}