---
library_name: mouse-core
tags:
- mouse-core
- reinforcement-learning
---

# micahr234/mouse-example-model-offline

This repository contains a MOUSE model checkpoint.

## Architecture

- Backbone: `qwen3`
- Hidden dimension: `1024`
- Heads: `action_value`
- Action head: `action_value`

### Encoder

`NumericEmbedder` reads flat step-record dicts and projects each declared modality
into the shared `1024`-dimensional token space before the
backbone.

| Field | Type | Required | Tensor shape | Dtype | Notes |
|---|---|---:|---|---|---|
| `action` | `discrete` | yes | `[B, S]` | `torch.long` | integer ids in `[0, 3]` |
| `observation` | `discrete` | yes | `[B, S]` | `torch.long` | integer ids in `[0, 63]` |
| `reward` | `fourier` | yes | `[B, S]` | `torch.float32` | scalar value |
| `done` | `discrete` | yes | `[B, S]` | `torch.long` | integer ids in `[0, 4]` |

## Install MouseCore

```bash
pip install mouse-core
```

## Load The Model

```python
import torch
from mouse_core import load_model
from mouse_core.models import preferred_dtype

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = (
    load_model("micahr234/mouse-example-model-offline", map_location="cpu")
    .eval()
    .to(device=device, dtype=preferred_dtype(device))
)
```

## Run Inference

Training and inference both take a `TokenBatch`. Training typically uses
`DataLoader(preparer=encoder.make_preparer())`. Online / inference builds one
with `encoder.prepare(rows)` where `rows` is `[B][S]` step-record dicts whose
keys match the encoder's declared modalities (and any `extra_fields`).

```python
# Batch shape: [B=1][S=1] — one sequence of one step.
batch = [[
    {
    "action": 0,
    "observation": 0,
    "reward": 0.0,
    "done": 0,
    }
]]
predictions, objective_data, cache = model(model.encoder.prepare(batch))

with torch.no_grad():
    predictions, _, cache = model(model.encoder.prepare(batch))
    action = model.get_action(predictions, temperature=0.0)
```

`model()` returns `(predictions, objective_data, cache)`. `objective_data` is a
`TensorDict[B, S]` of the modality tensors extracted by the encoder — pass it
to objectives during training. For cached incremental rollout, keep `cache` and
pass it back on the next call with `use_cache=True`. Cached batch rows may have
different lengths on every call (e.g. envs emitting different numbers of steps
between model calls): decoding runs through a FlexAttention session carried in
the cache, so each row decodes exactly as it would alone.