Metadata-Version: 2.4
Name: mouse-experiment
Version: 0.1.0
Summary: Meta-Optimization Using Sequential Experiences — experiment runner
License-Expression: GPL-3.0-or-later
Project-URL: Homepage, https://github.com/micahr234/mouse-experiment
Project-URL: Documentation, https://github.com/micahr234/mouse-experiment/tree/main/docs
Project-URL: Repository, https://github.com/micahr234/mouse-experiment
Project-URL: Issues, https://github.com/micahr234/mouse-experiment/issues
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: mouse-core>=0.1.0
Requires-Dist: mouse-env>=0.1.0
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: omegaconf>=2.3.0
Requires-Dist: tqdm>=4.67.3
Requires-Dist: huggingface_hub>=1.5.0
Requires-Dist: multiprocess>=0.70.18
Provides-Extra: tracking
Requires-Dist: trackio>=0.3.0; extra == "tracking"
Requires-Dist: wandb>=0.25.0; extra == "tracking"
Provides-Extra: dev
Requires-Dist: pyright; extra == "dev"
Requires-Dist: pytest; extra == "dev"
Requires-Dist: build>=1.2; extra == "dev"
Requires-Dist: twine>=6.1; extra == "dev"
Dynamic: license-file

# Meta-Optimization Using Sequential Experiences

<p align="center"><img src="https://raw.githubusercontent.com/micahr234/mouse-experiment/main/docs/mouse-experiment.png" width="400"/></p>

> **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](https://github.com/micahr234/mouse-core), 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](https://github.com/gradio-app/trackio)**.

For the core ML library, see **[mouse-core](https://github.com/micahr234/mouse-core)**.
For NS-Gym environments, see **[mouse-env](https://github.com/micahr234/mouse-env)**.

## Install

```bash
pip install mouse-experiment
```

Development:

```bash
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`](configs/ns_gym/default.yaml).

```bash
# 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

```yaml
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`](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](CONTRIBUTING.md).

## License

GNU General Public License v3.0 — see [LICENSE](LICENSE).
