Upload MOUSE model

README.md

@@ -0,0 +1,75 @@
+---
+library_name: mouse-core
+tags:
+- mouse-core
+- reinforcement-learning
+---
+
+# micahr234/mouse-example-model-augmented1
+
+This repository contains a MOUSE model checkpoint.
+
+## Architecture
+
+- Backbone: `qwen3`
+- Hidden dimension: `1024`
+- Heads: `action_value`
+- Action head: `action_value`
+
+### Encoder
+
+`StepEmbedder` 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]` |
+| - | `learnable` | no | `not read from step_stream` | `n/a` | learned tokens; no input field |
+
+## Install MouseCore
+
+```bash
+pip install mouse-core
+```
+
+## Load The Model
+
+```python
+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-augmented1", 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.
+
+```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(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`.

config.json

@@ -0,0 +1,98 @@
+{
+  "backbone": {
+    "hidden_dim": 1024,
+    "kwargs": {
+      "attention_bias": false,
+      "head_dim": 128,
+      "intermediate_size": 3072,
+      "max_position_embeddings": 40960,
+      "num_heads": 16,
+      "num_key_value_heads": 8,
+      "num_layers": 28,
+      "rms_norm_eps": 1e-06,
+      "use_sliding_window": false
+    },
+    "type": "qwen3"
+  },
+  "encoder": {
+    "kwargs": {
+      "concat_modalities": false,
+      "fourier_max": 10.0,
+      "fourier_min": 0.01,
+      "hidden_dim": 1024,
+      "include_type_token": false,
+      "modalities": [
+        {
+          "allow_none": false,
+          "field": "action",
+          "method": "rff",
+          "required": true,
+          "std": 0.02,
+          "tokens": 1,
+          "type": "discrete",
+          "vocab_size": 4
+        },
+        {
+          "allow_none": false,
+          "field": "observation",
+          "method": "rff",
+          "required": true,
+          "std": 0.02,
+          "tokens": 1,
+          "type": "discrete",
+          "vocab_size": 64
+        },
+        {
+          "allow_none": false,
+          "field": "reward",
+          "in_max": 100.0,
+          "in_min": 0.01,
+          "method": "rff",
+          "required": true,
+          "std": 0.02,
+          "tokens": 1,
+          "type": "rff"
+        },
+        {
+          "allow_none": false,
+          "field": "done",
+          "method": "rff",
+          "required": true,
+          "std": 0.02,
+          "tokens": 1,
+          "type": "discrete",
+          "vocab_size": 5
+        },
+        {
+          "allow_none": false,
+          "method": "rff",
+          "required": false,
+          "std": 0.02,
+          "tokens": 1,
+          "type": "learnable"
+        }
+      ],
+      "std": 0.02,
+      "token_data_len": 1,
+      "type_embedding_std": 0.0
+    },
+    "type": "step"
+  },
+  "format": "mouse-core-model-v1",
+  "heads": {
+    "action_head": "action_value",
+    "heads": [
+      {
+        "hidden_dim": 1024,
+        "in_features": 1024,
+        "name": "action_value",
+        "num_layers": 1,
+        "out_features": 4,
+        "scale": 0.1,
+        "type": "action_value",
+        "use_norm": true
+      }
+    ]
+  },
+  "hidden_dim": 1024
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08ec32b98075f2c9ea88b93e276040d08ee923e8022e0b00a625199c5054e3ac
+size 1762623279