File size: 2,528 Bytes
76bea55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | ---
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.
|