Views
No views yet
qwen31024actionactionStepEmbedder 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 | rff | yes | [B, S] | torch.float32 | scalar value |
done | discrete | yes | [B, S] | torch.long | integer ids in [0, 4] |
pip install mouse-core1import torch
2from mouse_core import load_model
3
4device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5model = load_model("micahr234/mouse-example-model-distill-online", map_location="cpu").eval().to(device)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.1# Batch shape: [B=1][S=1] — one sequence of one step.
2batch = [[
3 {
4 "action": 0,
5 "observation": 0,
6 "reward": 0.0,
7 "done": 0,
8 }
9]]
10predictions, objective_data, cache = model(batch)
11
12with torch.no_grad():
13 predictions, _, cache = model(batch)
14 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.