mouse-experiment 0.1.0

Meta-Optimization Using Sequential Experiences — experiment runner

Meta-Optimization Using Sequential Experiences

Warning: MOUSE is in early development and is not yet ready for use. APIs will change without notice.

mouse-experiment is the experiment runner for MOUSE, a modular PyTorch stack for in-context reinforcement learning. It provides YAML-driven configuration, offline training, online test rollouts, and logging to both Weights & Biases and Trackio.

For the core ML library, see mouse-core. For NS-Gym environments, see mouse-env.

Install

pip install mouse-experiment

Development:

git clone https://github.com/micahr234/mouse-experiment.git
cd mouse-experiment
source scripts/install.sh

Quick start

Each experiment is described by a single YAML file. A fully-annotated reference config is at configs/ns_gym/default.yaml.

# Activate environment
source .venv/bin/activate

# Run an experiment (full path or relative to configs/)
mouse-run configs/ns_gym/cartpole/train_ns_cartpole_with_bb.yaml

# Or via the tmux wrapper
./scripts/run.sh configs/ns_gym/cartpole/train_ns_cartpole_with_bb.yaml

Set the following environment variables before running:

Variable	Required	Purpose
HF_TOKEN	Yes	Hugging Face token for dataset and model push/pull
WANDB_TOKEN	Yes	Weights & Biases API key
TRACKIO_SPACE_ID	No	HF Space to sync Trackio dashboard to (e.g. user/my-space)

Experiment YAML structure

seed: 42

# ── Datasets ──────────────────────────────────────────────────────────────
load_train_dataset:
  name: play_ns_cartpole_big   # HF dataset id or local path
  split: train

load_eval_datasets:
  - name: play_ns_cartpole_big
    split: eval
    tag: id

save_dataset:
  name: my_rollout_dataset     # HF dataset to push rollouts to (null = skip)

# ── Model ─────────────────────────────────────────────────────────────────
model:
  backbone:
    type: llama                # llama | qwen3 | none
    load_config: meta-llama/Llama-3.2-1B
    load_pretrained_name: meta-llama/Llama-3.2-1B
    num_layers: 4
  vec_dqn_head:
    vec_dim: 2

# ── Training loop ─────────────────────────────────────────────────────────
loop:
  num_steps: 50000
  train_interval: 1
  eval_interval: 100
  test_interval: 1000
  lr: 1.0e-5

# ── Online test environments ───────────────────────────────────────────────
test_envs:
  cartpole_ns:
    id: NS-CartPole-v1
    num_envs: 10
    num_steps: 500
    action_source:
      name: learned_vec_dqn

# ── Logging ───────────────────────────────────────────────────────────────
wandb:
  project: "Train NS-Gym"
  run_name: null               # auto-generated from config filename if null

trackio:
  project: "Train NS-Gym"      # omit or set null to disable Trackio
  space_id: null               # e.g. "user/mouse-dashboard"

See configs/ns_gym/default.yaml for every available setting with inline documentation.

Experiment modes

Mode	train_interval	What it does
Play (collect data)	0	Runs test rollouts with expert q_star actions, saves trajectories to HF Hub
Train	> 0	Loads a prior play dataset, trains offline, runs periodic eval + online rollouts

Logging

Every run logs to:

• Weights & Biases — always. Project is set via wandb.project in the config.
• Trackio — when trackio.project is set. Metrics are stored locally in SQLite and optionally synced to a Hugging Face Space (via trackio.space_id or the TRACKIO_SPACE_ID env var).

Pre-built configs

Config	Environment	Mode
configs/ns_gym/cartpole/play_ns_cartpole.yaml	NS-CartPole-v1	play
configs/ns_gym/cartpole/train_ns_cartpole_with_bb.yaml	NS-CartPole-v1	train (Llama backbone)
configs/ns_gym/cartpole/train_ns_cartpole_without_bb.yaml	NS-CartPole-v1	train (no backbone)
configs/ns_gym/frozenlake/train_ns_frozenlake_with_bb.yaml	NS-FrozenLake-v1	train (Llama backbone)

Contributing

See CONTRIBUTING.md.

License

GNU General Public License v3.0 — see LICENSE.