Metadata-Version: 2.4
Name: mouse-env
Version: 0.3.0
Summary: Meta-Optimization Using Sequential Experiences — environments
License-Expression: GPL-3.0-or-later
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: gymnasium[atari,box2d]>=1.2.3
Requires-Dist: ns_gym>=1.0.10
Requires-Dist: stable-baselines3>=2.0.0
Requires-Dist: huggingface_hub>=1.5.0
Requires-Dist: pillow>=12.1.1
Requires-Dist: multiprocess>=0.70.18
Requires-Dist: tensordict>=0.7.0
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"
Requires-Dist: jupyter>=1.1.1; extra == "dev"
Requires-Dist: ipykernel>=6.29.5; extra == "dev"
Dynamic: license-file

# MOUSE Environments

<p align="center"><img src="docs/mouse-env.png" width="400"/></p>

> **Warning:** MOUSE is in early development and is not yet ready for production use. APIs may change without notice.

**mouse-env** is the environment layer for [MOUSE](https://github.com/micahr234/mouse-core), a modular PyTorch library for in-context reinforcement learning.

## Why mouse-env exists

In standard RL, episodes are independent — the algorithm resets, collects one episode, and repeats. What happened in episode 1 has no bearing on episode 2.

In **in-context RL**, the policy is a sequence model whose context window spans **multiple episodes**. The model adapts its behavior from recent history in a single forward pass, without gradient updates between episodes.

mouse-env is built for this regime. You call `step()` in a loop forever — the environment resets itself when an episode ends and keeps going. Episode boundaries show up as `done=1` or `done=2` in the data, and the following step delivers the fresh starting observation. Your training loop never needs to detect episode endings or call `reset()` itself.

## Core API

Using mouse-env looks like this:

- Build one vectorized environment with `make_vector_env(...)`
- Call `step(actions)` in a loop; receive `data`, `metadata`, and `metrics` every step

For each sub-environment index `i`:

- **`data[i]`** (model-visible)
  - `observation`
  - `reward`
  - `done`
  - `time`
- **`metadata[i]`** (training-only)
  - `q_star`
  - `reward_episodic`
  - `episode_index`
  - `group_id`
  - plus optional env-specific metadata (for example `ns_params`)
- **`metrics[i]`** (logging)
  - true episodic return and episode length, emitted only on episode end

This API is shared across tabular envs, Gymnasium control tasks, Atari, and non-stationary setups.

## Install

```bash
pip install mouse-env
```

Development setup:

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

## Quick start

```python
from mouse.envs import EnvConfig, make_vector_env

cfg = EnvConfig.cartpole(seed=0, num_envs=4, max_episode_steps=500)
env = make_vector_env(cfg)

for _ in range(1000):
    actions = env.sample_random_actions()
    data, metadata, metrics = env.step(actions)

env.close()
```

### Important stepping semantics

- **Use `step()` only — do not call `reset()`.**
  - The first `step()` performs an internal reset and returns the initial observation with `reward=0`, `done=0`, `time=0`.
  - Actions passed on that first call are ignored.
- **Autoreset is built in.**
  - After termination/truncation, the next emitted transition includes the reset observation (`reward=0`, `done=0`, `time=0`) before normal stepping resumes.
- **Observations/actions are typed dicts of tensors.**
  - Observation channels can include `discrete`, `continuous`, and/or `image`.
  - Actions follow the same keyed structure.
- **`done` is integer-coded.**
  - `0` = running, `1` = terminated, `2` = truncated.

See **[docs/guide.md](docs/guide.md)** for full field-level documentation, plus runnable notebooks in [`examples/`](examples/).

## Included environments and integrations

### 1) Procedural Frozen Lake

- **ID:** `Procedural-FrozenLake-v1`
- **Config helper:** `EnvConfig.procedural_frozenlake()`
- Random valid grid generation (size, holes, start/goal), optional per-goal rewards.
- Distinct from Gymnasium's fixed `FrozenLake-v1` benchmark.

### 2) Synthetic Environment

- **ID:** `SyntheticEnv-v1`
- **Config helper:** `EnvConfig.synthetic()`
- Random finite discrete MDP for controlled tabular experiments.

### 3) NS-Gym integration (external framework)

mouse-env integrates [NS-Gym](https://github.com/scope-lab-vu/ns_gym) to support non-stationary dynamics through the same API.

- `EnvConfig.ns_cartpole(non_stationary_params={...})` for scheduler/update config
- `NSGymInterfaceWrapper` to normalize observations + expose `metadata[i]["ns_params"]`
- Vectorized non-stationary streams with standard `(data, metadata, metrics)` outputs

Example: [examples/03_ns_gym_oscillating.ipynb](examples/03_ns_gym_oscillating.ipynb)  
NS-Gym docs: [nsgym.io](https://nsgym.io/)

### 4) Atari integration

mouse-env keeps ALE/Gymnasium Atari semantics intact and exposes them through the same API.

- `EnvConfig.atari()` preset:
  - `AtariPreprocessing`
  - frame skip 4
  - grayscale 84×84
  - noop warm-up
- Frames are surfaced in `data[i]["observation"]["image"]` (flattened)

Requirement: `gymnasium[atari]` (`ale_py`).  
Example: [examples/04_atari_preprocessing.ipynb](examples/04_atari_preprocessing.ipynb)

## Training-oriented features

### Expert Q-values (`q_star_source`)

- Exposed as `metadata[i]["q_star"]`.
- Backends:
  - **`sb3_rl_zoo`**: pretrained Stable-Baselines3 policies from Hugging Face Hub (CartPole preset uses this by default).
  - **`metadata_q_star`**: exact Q* solved from tabular MDP dynamics (Procedural Frozen Lake + Synthetic Environment).

### Partial observability

Use `observation_indices` to mask dimensions on continuous-vector observation spaces.  
Example: [examples/05_partial_observability.ipynb](examples/05_partial_observability.ipynb)

### Reward shaping

Use `reward_scale` and `reward_shift`; normalized training signal appears in `metadata[i]["reward_episodic"]`.  
Example: [examples/06_reward_shaping.ipynb](examples/06_reward_shaping.ipynb)

## Contributing

See [CONTRIBUTING.md](CONTRIBUTING.md).

## License

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