TSO Sensorium

A library for sensor data recording, post-processing, and episodic dataset creation for robot learning.

Package layout

tso_sensorium/
├── recording/     # sensor capture: core + ros1/ (Noetic) and ros2/ (Jazzy) adapters
├── processing/    # stereo rectification, deinterlacing, timestamp alignment
├── episodes/      # schema, episode assembly, parallel dataset building
├── export/        # dataset writers: CSV folders, LeRobot (v3.0 format)
└── scripts/       # command-line entry points

Processing, episode assembly, and export have no ROS dependency and run anywhere. Recording adapters import the ROS client libraries and are only usable inside the corresponding ROS environment.

Installation

With pixi (recommended)

The preferred workflow is to enter the environment once and run commands normally inside it. Pixi creates the environment on first use and installs this package in editable mode automatically. Install Pixi first if it is not already available.

For ROS 1 development, start a clean terminal and enter the ROS 1 environment:

pixi shell -e ros1

Confirm that Python is using the RoboStack ROS installation from this repository's Pixi environment:

python -c "import rospy; print(rospy.__file__)"

The path must be under .pixi/envs/ros1/, not /opt/ros/noetic/. Once inside the shell, run ROS commands, application commands, and tests directly.

The tso_testbed configurations also require the custom messages generated by the testbed catkin workspace. Add only that workspace's generated Python packages to the Pixi environment:

export TESTBED_WORKSPACE=/path/to/robot_testbed
export PYTHONPATH="$TESTBED_WORKSPACE/devel/lib/python3/dist-packages"
python -c "import rospy, testbed_msgs; print(rospy.__file__); print(testbed_msgs.__file__)"

rospy must still resolve under .pixi/envs/ros1/, while testbed_msgs must resolve under the testbed workspace. Do not source the workspace's setup.bash inside the Pixi shell: it also adds /opt/ros/noetic and replaces RoboStack's ROS Python packages with the system installation.

The following commands can then use the real testbed configuration:

python -m tso_sensorium.scripts.record \
    --config_path configs/recording/tso_testbed.yaml \
    --output_folder "$HOME/tso_sensorium_recordings"
pytest tests/recording/test_ros1_adapter.py

Use a separate terminal for each foreground process and run pixi shell -e ros1 in each one. Leave an environment with exit or Ctrl-D. The other environments work the same way:

pixi shell             # core development and non-ROS tests
pixi shell -e ros2     # ROS 2 recording and tests

Do not automatically source /opt/ros/noetic/setup.bash or a catkin workspace from .bashrc. Those scripts prepend the complete system ROS installation to PYTHONPATH, which can make Pixi's Python import the wrong ROS packages. If the current terminal has already sourced system ROS, open a clean terminal before entering the Pixi shell, then add only the generated message path as shown above.

For scripts and CI, where an interactive shell is inconvenient, use the equivalent one-shot form:

pixi run test
pixi run -e ros1 pytest tests/recording/test_ros1_adapter.py
pixi run -e ros2 pytest tests/recording/test_ros2_adapter.py

The environments are declared in pyproject.toml and locked in pixi.lock. On machines with a small home quota, set PIXI_CACHE_DIR to storage with more space before creating an environment. LeRobot export remains a pip extra to avoid duplicating a multi-gigabyte torch installation per checkout.

Alternative installations

Pip

pip install -e .              # core: processing, episodes, export
pip install -e ".[gui]"       # + browser dashboard (Flask)
pip install -e ".[lerobot]"   # + LeRobot dataset export (Python >= 3.10)
pip install -e ".[test]"      # + pytest

Python requirements: 3.9+ (the pydantic-backed config layer); LeRobot export additionally needs 3.10+. Extras combine as usual, e.g. pip install -e ".[gui,test]". After install the scripts are available as console commands (tso-record, tso-record-ui, tso-annotate, tso-generate-dataset, tso-label-phases, and the *-ros2 variants), equivalent to the python -m tso_sensorium.scripts.* forms used below.

Mamba for ROS recording

For recording without Pixi, ROS client libraries must come from a system ROS installation or a dedicated RoboStack environment because rospy and rclpy are not available from PyPI. The environment files under environments/ provide the latter:

mamba env create -f environments/ros1-noetic.yml   # or ros2-jazzy.yml
mamba activate tso-sensorium-ros1
pip install -e . --no-deps                          # deps already come from conda

Note: the Noetic robot PC runs Python 3.8 and the LeRobot exporter needs Python 3.10+, so recording and LeRobot export are not meant to share one environment. Record on the robot, export on the processing machine.

User guide

Configuration

Every script is driven by a pydantic-validated YAML file under configs/, selected with --config_path. The same rules apply everywhere:

• Unknown keys fail loudly. A typo'd key aborts with a validation error naming the key, instead of being silently ignored.
• Polymorphic entries pick their variant with a type: key — recorder kinds (topic, video), frame transforms (deinterlace, resize, rectify), table transforms (parse_vector3, sum_columns, ...), writers (csv, lerobot), phase labelers and triggers. The fields next to type are the variant's own options.
• Configs compose with !include, resolved relative to the including file. This keeps shared pieces in one place, e.g. configs/recording/mock_ui.yaml includes configs/dataset/mock.yaml as its generation: section, and the same file is included by configs/annotation/mock.yaml.
• Calibration files packaged with the library can be referenced as package://<asset name>.

CLI overrides: any field can be overridden with --dot.path value pairs on top of the YAML. Values are parsed as YAML scalars, so true, 3.5, and [a, b] become their typed equivalents. List elements are addressed by index, and setting a type key resets that section to the new variant (follow-up flags then fill its fields):

python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings \
    --save_frames true \
    --videos.0.frame_column framePath \
    --writer.type lerobot --writer.output_root /data/lerobot

All modules require Python 3.9+ (pydantic-backed config layer).

Recording

When the endoscope is mounted on a moving robot, publish its calibrated live camera orientation from the robot-state topic in a separate ROS 1 Pixi shell:

python -m tso_sensorium.scripts.publish_robot_camera_transform \
    --subscribe /ur5e_rcm_twist_controller/RobotState \
    --topic /robot_camera_transform \
    --calibration_path configs/calibration/tso_endoscope_mount.yaml

The calibration file contains the fixed EE-to-camera axis mapping and optical tilt. Replace it for a different holder or endoscope geometry. The equivalent installed command is tso-publish-camera-transform.

Record ROS topics and video streams in one shot (inside a ROS 1 environment):

export TESTBED_WORKSPACE=/path/to/robot_testbed
export PYTHONPATH="$TESTBED_WORKSPACE/devel/lib/python3/dist-packages"
python -m tso_sensorium.scripts.record \
    --config_path configs/recording/tso_testbed.yaml \
    --output_folder "$HOME/tso_sensorium_recordings"

For interactive sessions, record_ui runs the browser UI (requires the gui extra):

export TESTBED_WORKSPACE=/path/to/robot_testbed
export PYTHONPATH="$TESTBED_WORKSPACE/devel/lib/python3/dist-packages"
python -m tso_sensorium.scripts.record_ui \
    --config_path configs/recording/tso_testbed_ui.yaml \
    --session.output_folder "$HOME/tso_sensorium_recordings"

For the force-sensing setup, the bundled UI config records UR5e state, the left, right, and combined endoscope streams, the robot-camera transform, and the Bota wrench and IMU topics. It saves under ~/tso_sensorium_recordings by default:

python -m tso_sensorium.scripts.record_ui \
    --config_path configs/recording/force_session_ui.yaml

Video readiness uses an explicitly configured lightweight per-frame status topic, not raw image subscriptions. Each video entry in the recording YAML must specify liveness_topic when used with the UI; topic names are never inferred. The bundled setups use sensor_msgs/CameraInfo messages. Recording still subscribes to the configured image topic when an episode starts and releases that subscription when it stops. The configured live preview is a separate, intentional image subscription.

Set the status topic and its message type alongside the recorded image topic:

- type: video
  file_name: camera
  topic_name: /camera/image_raw
  liveness_topic: /camera/camera_info
  liveness_message_type: sensor_msgs.msg.CameraInfo

Other per-frame status messages, such as std_msgs.msg.Header, can be selected with liveness_message_type. Command-line recording does not use these fields. The status topic must publish with each acquired frame; latched calibration alone does not establish ongoing camera freshness. A missing or stale status topic keeps the source unready rather than silently falling back to expensive raw-image subscriptions. If a camera can only provide raw images for status, explicitly select that image topic and sensor_msgs.msg.Image, accepting its bandwidth cost. This affects readiness only, not the recording format.

Open http://<host>:8080 from any machine on the network. The dashboard has two sections: Record (live camera feed, per-sensor liveness, start/stop of demonstrations, per-episode topic selection) and Library (episode browsing with in-browser replay, phase annotation, and dataset generation — see below). The server is unauthenticated; expose it on trusted networks only.

To try the full record → browse → annotate → generate loop without any hardware, first enter pixi shell -e ros1, then run the bundled synthetic sensors:

First check whether the current ROS_MASTER_URI already has a running master:

rostopic list

If that succeeds, reuse the existing master. If it fails, start roscore in that terminal. Do not start a second master:

roscore

Open a second terminal, enter pixi shell -e ros1, and start the sensors:

python -m tso_sensorium.scripts.mock_sensors

Open a third terminal, enter pixi shell -e ros1, and start the dashboard:

python -m tso_sensorium.scripts.record_ui \
    --config_path configs/recording/mock_ui.yaml

ROS may warn when ~/.ros/log exceeds 1 GB. Inspect it with rosclean check; run rosclean purge only if deleting old ROS logs is acceptable.

The same dashboard runs on ROS 2 (record_ui_ros2, mock_sensors_ros2) with identical configs — recorder entries reference message types by dotted path (std_msgs.msg.Bool), which resolve to the ROS 2 classes inside a ROS 2 environment:

python -m tso_sensorium.scripts.mock_sensors_ros2 &
python -m tso_sensorium.scripts.record_ui_ros2 \
    --config_path configs/recording/mock_ui.yaml

Annotation and dataset studio (no ROS required)

The Library section also runs as a standalone app on any machine — a processing workstation without ROS, pointed at a folder of recordings:

python -m tso_sensorium.scripts.annotate \
    --config_path configs/annotation/mock.yaml
# or directly:
python -m tso_sensorium.scripts.annotate \
    --recordings_root /data/recordings --port 8090

From the dashboard you can:

• Switch the recordings folder being browsed (the path field at the top; the episode list, annotations, and legend reload for that folder).
• Edit the dataset metadata and phase legend: dataset name, task, and the mapping of integer phase labels to a phase name plus its language instruction variants. One instruction per phase is deterministic; several lines make a stochastic mapping — one variant is sampled per episode, seeded by the episode name so regeneration is reproducible. The legend is saved to dataset_metadata.json at the recordings root.
• Annotate episodes on a timeline (pencil icon next to an episode): colored phase segments under the video, click to seek and select, drag segment edges, split at the playhead, and assign phases from the legend. Edits are saved to annotations.json inside the episode folder with source: manual, so re-running an automatic labeler never overwrites them.
• Generate datasets with per-run options (format, output root, frame extraction, sync tolerance); discarded episodes are listed with the reason.

Phase legends

Legends map phase labels to names and instruction variants. Select one in your dataset-generation config:

annotations:
  legend: package://instructions/endoscope_guidance.yaml
  legend_source: config
  language_source: phase_legend

legend also accepts a YAML filename or an inline metadata mapping.

• legend_source: config uses the supplied legend; auto (default) prefers the phase legend in the recordings root's dataset_metadata.json, falling back to the configured legend when none is saved.
• language_source: phase_legend samples instructions from the legend; annotation (default) uses segment-specific text when present, otherwise falling back to the legend.

See the annotation configuration for all fields and defaults.

Automatic phase labeling

Labelers segment episodes from recorded signals and write automatic segments into each episode's annotations.json:

python -m tso_sensorium.scripts.label_phases \
    --config_path configs/labeling/bowel_retraction.yaml \
    --recordings_root /data/recordings

configs/labeling/ shows both labeler kinds: column_threshold (two phases split by a signal threshold) and sequential_trigger (an ordered phase sequence where each trigger hands over to the next phase — gripper state changes, motion starting or settling). The typical workflow is auto-label → correct visually in the dashboard → generate.

Dataset generation

python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings

# Same recordings, exported as a LeRobot dataset instead:
python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings \
    --writer.type lerobot --writer.output_root /data/lerobot

# Recompute transformed actions without re-encoding observations or videos:
python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings \
    --writer.type lerobot_action_update \
    --writer.dataset_root /data/lerobot

When the config has an annotations: section, each episode's phase segments are joined onto the aligned table as an integer phase column and a language instruction column. Every writer stores generation statistics and any configured schema.coordinate_frame_features in dataset_metadata.json. Coordinate-frame entries list their exact component columns, named frame, and whether that frame is fixed, moving, or unknown across timesteps. LeRobot exports can also carry named schema.auxiliary_features, such as an integer phase label, without folding them into the policy action vector. Action-only updates require an existing LeRobot v3 dataset whose episode order, lengths, schema, observations, auxiliary values, and task strings exactly match the regenerated episodes. They atomically replace only actions, derived action statistics, and dataset_metadata.json; a validation failure or cancellation leaves the existing dataset unchanged.

Library usage

Generate an episode table programmatically:

from tso_sensorium.episodes import EpisodeGenerator

episode = (
    EpisodeGenerator()
    .add_video(
        video_path="episode/left.mp4",
        timestamps_path="episode/left_timestamps.csv",
        sync_col_name="timestamp",
        frames_output_path="episode/frames/left",
        frame_col_name="left_frame",
        save_frames=True,
    )
    .add_state(
        state_data_path="episode/robot_state.csv",
        sync_col_name="timestamp",
        dataset_cols=["x", "y", "z", "roll"],
    )
    .generate_dataset()
    .save_dataset("episode/episode.csv")
)

Rectify and deinterlace stereo images:

from tso_sensorium.processing import Rectifier, deinterlace_cv_image

rectifier = Rectifier(calibration_file_path="calibration.yml")
left, right = deinterlace_cv_image(image=interlaced_image)
left, right = rectifier.rectify(left=left, right=right)

Development

Enter the default Pixi shell for framework-independent development:

pixi shell
pytest
pytest -m ""
ruff format tso_sensorium/ tests/
ruff check tso_sensorium/ tests/

ROS adapter tests require the matching environment and skip in the default environment:

exit
pixi shell -e ros1
pytest tests/recording/test_ros1_adapter.py

For non-interactive use, the equivalents are pixi run test, pixi run test-all, pixi run lint, and pixi run -e ros1 pytest tests/recording/test_ros1_adapter.py.

CI checks pull requests and main, including ROS 1/2 tests and package builds. Published GitHub releases upload to PyPI after those checks pass.