File size: 3,097 Bytes
adaf847 | 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 84 | #!/usr/bin/env python3
"""TinyLunarPolicy - a 14-parameter neural network that lands Gymnasium LunarLander-v3.
logits = obs @ W # obs: (8,), W: (8, 4)
action = argmax(logits) # action in {0,1,2,3}
W is not free: it is built from **14 trainable scalars** by hard-coding the exact
mirror symmetry of the environment (x -> -x, vx -> -vx, angle -> -angle, w -> -w,
legs swap, action left <-> right; actions 0 "nothing" and 2 "main engine" are
self-paired). So the model has 14 parameters.
Dependencies: torch, safetensors, numpy.
"""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
OBS_DIM = 8
N_ACTIONS = 4
N_PARAMS = 14
# Gymnasium LunarLander-v3 observation layout (for reference).
OBS_NAMES = [
"x", "y", "vx", "vy", "angle", "angular_velocity",
"leg_left_contact", "leg_right_contact",
]
class TinyLunarPolicy(nn.Module):
"""Linear policy 8 -> 4 built from 14 parameters via the mirror symmetry."""
def __init__(self) -> None:
super().__init__()
self.sym_theta = nn.Parameter(torch.zeros(N_PARAMS))
@property
def n_params(self) -> int:
return int(self.sym_theta.numel())
def weights(self) -> torch.Tensor:
"""Materialise W (8, 4) from the 14 parameters; logits = obs @ W."""
f = self.sym_theta
W = torch.zeros(OBS_DIM, N_ACTIONS, dtype=f.dtype)
# actions 0 and 2 are self-paired under the mirror: only y, vy and leg contact survive
W[1, 0], W[3, 0], W[6, 0], W[7, 0] = f[0], f[1], f[2], f[2]
W[1, 2], W[3, 2], W[6, 2], W[7, 2] = f[3], f[4], f[5], f[5]
# actions 1 and 3 form the chiral pair: W[:, 3] = M @ W[:, 1]
W[:, 1] = f[6:14]
W[0, 3], W[1, 3], W[2, 3], W[3, 3] = -W[0, 1], W[1, 1], -W[2, 1], W[3, 1]
W[4, 3], W[5, 3], W[6, 3], W[7, 3] = -W[4, 1], -W[5, 1], W[7, 1], W[6, 1]
return W
def numpy_weights(self) -> np.ndarray:
return self.weights().detach().cpu().numpy()
def forward(self, obs: torch.Tensor) -> torch.Tensor:
"""obs: (..., 8) -> logits (..., 4)."""
return obs @ self.weights()
@torch.no_grad()
def act(self, obs) -> int:
o = torch.as_tensor(np.asarray(obs, dtype=np.float32))
return int(torch.argmax(self.forward(o)).item())
# -- loading -------------------------------------------------------------
@classmethod
def from_safetensors(cls, path: str) -> "TinyLunarPolicy":
from safetensors.torch import load_file
model = cls()
model.load_state_dict(load_file(path))
model.eval()
return model
@classmethod
def from_pretrained(cls, repo_id_or_path: str = "Dimitrius174/lunarlander-14params",
filename: str = "model.safetensors") -> "TinyLunarPolicy":
import os
if os.path.isdir(repo_id_or_path):
return cls.from_safetensors(os.path.join(repo_id_or_path, filename))
from huggingface_hub import hf_hub_download
return cls.from_safetensors(hf_hub_download(repo_id_or_path, filename))
|