micahr234/ns_gym_frozenlake_without_bb

MOUSE — Meta-Optimization Using Sequential Experiences

A context-conditioned sequence model for reinforcement learning. It reads a history of environment transitions and outputs action logits.

Install (requires private repo access)

pip install "git+https://github.com/micahr234/MOUSE.git"

Load

from mouse.models.base import load_model

model = load_model("micahr234/ns_gym_frozenlake_without_bb")
model.eval()

Step stream

The model takes a TensorDict[B, S] — B parallel sequences of S timesteps each. This model was trained with S = 100; keep context close to that.

import torch
from tensordict import TensorDict

B, S = 1, 1   # S grows each step when using the cache

step_stream = TensorDict(
    {
        "action":         torch.zeros(B, S, dtype=torch.int64),
        "reward":         torch.zeros(B, S, dtype=torch.float32),
        "done":           torch.zeros(B, S, dtype=torch.int64),  # 0=alive 1=terminal 2=truncated
        "obs_discrete":   torch.zeros(B, S, dtype=torch.int64),
    },
    batch_size=(B, S),
)

Inference

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

with torch.no_grad():
    out, cache = model(step_stream.to(device))

out is a TensorDict[B, S] with one key per enabled head (A = model.max_num_actions, D = vec_dim):

Key	Shape	Description
dqn	[B, S, A]	Q-value logits (online)
dqn_target	[B, S, A]	Q-value logits (target)
vec_dqn	[B, S, A, D]	Action vectors (online); use get_action or vec_dqn_scores
vec_dqn_target	[B, S, A, D]	Action vectors (target)

Select an action from the last timestep (works for all heads, handles temperature):

# greedy (temperature=0) or stochastic (temperature>0)
action = model.get_action(out, head="vec_dqn", temperature=0.0)  # [B]

Online rollouts with KV-cache

Pass one new step at a time (S=1) and carry the cache forward to avoid re-processing the full history on every call:

cache = None

while not done:
    step_stream = TensorDict(
        {
            "action":         last_action.unsqueeze(1),
            "reward":         last_reward.unsqueeze(1),
            "done":           last_done.unsqueeze(1),
            "obs_discrete":   obs_disc.unsqueeze(1),     # [B, 1]
        },
        batch_size=(B, 1),
    )

    with torch.no_grad():
        out, cache = model(step_stream.to(device), cache=cache, use_cache=True)
    action = model.get_action(out, head="vec_dqn", temperature=0.0)
    step_idx += 1

Cache warning. This model was trained on sequences of length 100. Quality degrades once the cache exceeds roughly 2× that length — reset it (cache = None) before that limit.