mouse-example-model-distill-online

micahr234/mouse-example-model-distill-online

This repository contains a MOUSE model checkpoint.

Architecture

• Backbone: qwen3
• Hidden dimension: 1024
• Heads: action
• Action head: action

Encoder

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

Install MouseCore

pip install mouse-core

Load The Model

import torch
from mouse_core import load_model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = load_model("micahr234/mouse-example-model-distill-online", map_location="cpu").eval().to(device)

Run Inference

The model accepts a list[list[dict]] batch of shape [B][S] — B sequences, each containing S step-record dicts with flat keys matching the encoder's declared modalities above.

# 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(batch)

with torch.no_grad():
    predictions, _, cache = model(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 one-step rollout, keep cache and pass it back on the next call with use_cache=True.