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
pip install mouse-core
Load The Model
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).
# 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.