File size: 3,595 Bytes
386302b bb8222a 386302b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | #!/usr/bin/env python3
"""TinyLunarPolicy — a 32-parameter single-layer neural network that lands LunarLander-v3.
logits = Linear(8 -> 4, bias=False)(obs) # 32 parameters
action = argmax(logits)
This is the *minimal* policy that solves Gymnasium `LunarLander-v3`:
mean return ~231.2, 99.7% landings (0.25% crashes), 60/60 evaluation seeds solved (see README).
It is exactly a single linear layer + argmax, and is mathematically identical to a
single-centroid nearest-neighbour policy in the C2 space (s (x) onehot(a)).
Dependencies: torch, safetensors (numpy optional, only for `.act`).
"""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
OBS_DIM = 8
N_ACTIONS = 4
# Gymnasium LunarLander-v3 observation layout (for interpretability only).
OBS_NAMES = [
"x", "y", "vx", "vy", "angle", "angular_velocity",
"leg_left_contact", "leg_right_contact",
]
class TinyLunarPolicy(nn.Module):
"""Linear 8->4 policy head; the action is the argmax of the four logits."""
def __init__(self, obs_dim: int = OBS_DIM, n_actions: int = N_ACTIONS,
bias: bool = False) -> None:
super().__init__()
self.obs_dim = int(obs_dim)
self.n_actions = int(n_actions)
self.linear = nn.Linear(self.obs_dim, self.n_actions, bias=bias)
# ── introspection ────────────────────────────────────────────────────────
@property
def n_params(self) -> int:
return sum(p.numel() for p in self.parameters())
def weights(self) -> np.ndarray:
"""Return W of shape (obs_dim, n_actions) such that logits = obs @ W."""
return self.linear.weight.detach().cpu().numpy().T # (8, 4)
# ── policy ───────────────────────────────────────────────────────────────
def forward(self, obs: torch.Tensor) -> torch.Tensor:
"""obs: (..., 8) -> logits (..., 4)."""
return self.linear(obs)
@torch.no_grad()
def logits(self, obs) -> torch.Tensor:
o = torch.as_tensor(np.asarray(obs, dtype=np.float32))
return self.forward(o)
@torch.no_grad()
def act(self, obs) -> int:
"""Greedy action for a single observation (array-like of length 8)."""
return int(torch.argmax(self.logits(obs)).item())
# ── loading ──────────────────────────────────────────────────────────────
@classmethod
def from_safetensors(cls, path: str, **kwargs) -> "TinyLunarPolicy":
"""Load from a local `.safetensors` file."""
from safetensors.torch import load_file
model = cls(**kwargs)
model.load_state_dict(load_file(path))
model.eval()
return model
@classmethod
def from_pretrained(cls, repo_id_or_path: str = "Dimitrius174/lunarlander-lin32",
filename: str = "model.safetensors", **kwargs) -> "TinyLunarPolicy":
"""Load from the Hugging Face Hub (or a local directory)."""
import os
if os.path.isdir(repo_id_or_path):
return cls.from_safetensors(os.path.join(repo_id_or_path, filename), **kwargs)
from huggingface_hub import hf_hub_download
return cls.from_safetensors(hf_hub_download(repo_id_or_path, filename), **kwargs)
|