WormRL Internal README

This repository trains and runs discrete-action SAC policies for closed-loop optogenetics on worm calcium-imaging data. The current codebase has two primary workflows and they must stay separate.

Current Contract

1. sac/mainline_replay.py is the offline model-development loop. It reads a complete H5 file directly, builds feature matrices directly from H5, does a time split, trains on the early frames, validates on the later frames, and writes model-test artifacts. It does not subscribe to ZMQ, publish ZMQ, or replay checkpoints.

2. sac/zmq_closed_loop.py is the actual worm closed-loop program. It trains/prepares from H5, then runs online inference from ZMQ frames. real mode subscribes to external microscope ZMQ. dummy mode self-publishes H5-derived fake ZMQ frames to test the subscriber, online projection, inference guard, and optional LabJack output.

3. Training does not need ZMQ. Training needs features: [T, neurons] and light: [T]. These come from H5 through sac/h5_to_feature.py.

4. ZMQ is a communication layer for online inference and dummy dry runs. sac/h5_to_zmq.py should not be part of the normal training path.

5. Use the pixi environment. Run commands as pixi run python ... unless there is a specific reason not to.

Recommended Offline Validation Data

Use this as the first default for offline model-parameter testing:

data/Raw_Data/20260316_immobile_OptoAct_AWCON/w4/data_loader_saved.h5

Why:

• 3605 frames, about 2.0 minutes at 30 fps.
• 113 neurons.
• green, red, and ratio have no missing feature pairs.
• 16 light events, about 5 frames each.
• The current 70/30 time split has similar light statistics:
  • train: light ratio about 2.18%, about 7.85 events/min, 11 events
  • test: light ratio about 2.31%, about 8.32 events/min, 5 events

Second choice:

data/Raw_Data/20260316_immobile_OptoAct_AWCON/w5/data_loader_saved.h5

It also has clean feature coverage, but the 30% validation segment has only 2 light events, so metrics are noisier.

Avoid using outputs/closed_loop_data/* as the main 70/30 validation set. Those files often have very different train/test light frequencies because they are closed-loop outputs. They are useful for stress tests or behavior analysis, not clean parameter validation.

Split Behavior

The mainline split is currently a direct chronological split:

split_idx = floor(total_frames * split_ratio)
train = frames[:split_idx]
test  = frames[split_idx:]

Default split_ratio is 0.70.

If train/test light frequency differs strongly, do not trust a single 70/30 score as a clean validation result. Compare train/test light ratio, event count, and events per minute. For unstable data, prefer a block-based or event-balanced validation design before drawing conclusions.

Main Commands

Offline Model Test

pixi run python -c "
from pathlib import Path
from sac.mainline_replay import run_mainline_split

run_mainline_split(
    h5_path=Path('data/Raw_Data/20260316_immobile_OptoAct_AWCON/w4/data_loader_saved.h5'),
    out_dir=Path('outputs/mainline_model_test'),
    split_ratio=0.70,
    signal_source='green',
    channel_index=1,
    device='cuda:0',
)
"

run_train_zmq_split() still exists only as a compatibility wrapper. New code should call run_mainline_split().

Closed-Loop Dummy Dry Run

pixi run python -c "
from pathlib import Path
from sac.zmq_closed_loop import run_online_workflow

run_online_workflow(
    h5_path=Path('data/Raw_Data/20260316_immobile_OptoAct_AWCON/w4/data_loader_saved.h5'),
    out_dir=Path('outputs/closed_loop_dummy'),
    mode='dummy',
    limit_frames=200,
    signal_source='green',
    channel_index=1,
    labjack_enabled=False,
)
"

Closed-Loop Real Mode

pixi run python -c "
from pathlib import Path
from sac.zmq_closed_loop import run_online_workflow

run_online_workflow(
    h5_path=Path('data/Raw_Data/20260316_immobile_OptoAct_AWCON/w4/data_loader_saved.h5'),
    out_dir=Path('outputs/closed_loop_real'),
    mode='real',
    zmq_host='127.0.0.1',
    zmq_port=33091,
    signal_source='green',
    channel_index=1,
    max_frames=1000,
    labjack_enabled=True,
)
"

Use labjack_enabled=True only when the hardware output should actually be active.

Syntax and Import Smoke Checks

pixi run python -m py_compile sac/h5_to_feature.py sac/mainline_replay.py sac/zmq_closed_loop.py
pixi run python -c "from sac.mainline_replay import run_mainline_split; from sac.zmq_closed_loop import run_online_workflow; from sac.h5_to_feature import load_h5_frame_data; print('imports ok')"

Data Flow

Mainline

H5
  -> h5_to_feature.load_h5_frame_data()
  -> features/light/frame_times/neuron_ids
  -> chronological split
  -> train_sac_from_arrays(train_features, train_light)
  -> PCA/state/replay/training inside train_sac.py
  -> project all features with train PCA
  -> predict validation frames

Closed-Loop Real

H5 -> features/light -> train SAC -> save/load policy artifacts
external microscope ZMQ -> OnlineFeatureProjector -> ensemble action -> guard -> optional LabJack

Closed-Loop Dummy

H5 -> features/light -> train SAC
H5-derived fake ZMQ messages -> local PUB/SUB -> OnlineFeatureProjector -> ensemble action -> optional LabJack

Module Ownership

• sac/h5_to_feature.py Shared H5 feature loading and small helpers. Its primary output is H5FrameData(features, light, frame_times, neuron_ids, coverage, light_intervals, meta). It intentionally stops before fitting PCA because PCA/state generation depends on the train split.

• sac/mainline_replay.py Offline split-mode model testing. Keep it H5-only.

• sac/zmq_closed_loop.py Hardware-facing closed-loop workflow. Keep real/dummy behavior stable and be cautious with LabJack changes.

• sac/train_sac.py Core SAC training. It fits PCA, builds stacked states, builds replay, trains the vectorized SAC ensemble, and writes artifacts.

• sac/data_utils.py PCA/state/replay utilities plus OnlineFeatureProjector for online ZMQ frames.

• sac/h5_to_zmq.py Utility for H5-derived fake ZMQ messages and optional publishing. Use it for dummy mode or ZMQ testing, not for normal training.

• sac/naive_baseline.py Rule-based comparison baseline. Some paths still use H5-derived fake ZMQ helpers.

Important Parameters

• signal_source H5 table family used as the canonical source for IDs/times: green, red, or ratio.

• channel_index Column inside ZMQ-style message fluorescence[:, channel_index]. In H5-direct loading, it selects the corresponding source from output_channel_sources.

• output_channel_sources Usually "red,green,ratio", so channel_index=1 selects green.

• pca_dim and stack_frames State dimension is pca_dim * stack_frames. Current defaults in TrainConfig are 10 and 3, so state dim is 30.

• fill_missing Missing-value policy is shared with the old H5-to-ZMQ-like path through h5_to_zmq.apply_fill_policy(). Keep this consistent when comparing runs.

Output Artifacts

Typical run directory contains:

• config.json
• dataset_stats.json
• pseudo_reward_reference.json
• target_state.npy
• pca_components.npy
• pca_mean.npy
• checkpoints/agent_XX.pt
• best_checkpoints.json
• train_log.jsonl
• agent_stop_summary.json
• ensemble_summary.json
• mainline-only validation CSV/HTML summaries when using run_mainline_split()

Development Rules

• Keep mainline_replay.py and zmq_closed_loop.py separate.
• Do not make zmq_closed_loop.py depend on high-level mainline workflows.
• Do not put argparse/click parsing inside core business logic.
• Core functions should accept plain parameters and stay importable.
• Use python -m module.path for module entry points.
• Read relevant files before making changes; keep edits minimal around hardware-facing code.
• Do not enable LabJack output by default in test commands.