File size: 1,792 Bytes
1f7a5ce | 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 | """Small permutation-equivariant policy baseline on the same current-state facts.
This is a standalone supervised MLP, not Qwen, Jev or a language-model adapter.
"""
import torch
from torch import nn
FEATURES=['lives/3','pellets_left/244','to_junction/10','frightened_seconds_left/6',
'ghost_distance/15','ghost_present','ghost_coming','edible_distance/15',
'edible_present','nearby_pellets/6','food_distance/40','food_present',
'power_distance/25','power_present','is_back']
def encode(state,keys):
global_features=[state['lives']/3,state['pellets_left']/244,state['to_junction']/10,
state.get('frightened_seconds_left',0)/6]
rows=[]
for key in keys:
f=state['options'][key]
dist=lambda field,scale: 1. if f.get(field) is None else f[field]/scale
rows.append(global_features+[dist('ghost',15),float(f.get('ghost') is not None),
float(f.get('ghost_coming',False)),dist('edible',15),float(f.get('edible') is not None),
f['pellets']/6,dist('food',40),float(f.get('food') is not None),dist('power',25),
float(f.get('power') is not None),float(key=='back')])
return rows
class StructuredPolicy(nn.Module):
def __init__(self):
super().__init__()
self.option=nn.Sequential(nn.Linear(15,64),nn.SiLU(),nn.Linear(64,64),nn.SiLU())
self.score=nn.Sequential(nn.Linear(192,64),nn.SiLU(),nn.Linear(64,1))
def forward(self,x,valid):
h=self.option(x)
mean=(h*valid[...,None]).sum(1)/valid.sum(1,keepdim=True)
peak=h.masked_fill(~valid[...,None],-1e9).max(1).values
pooled=torch.cat([h,mean[:,None,:].expand_as(h),peak[:,None,:].expand_as(h)],-1)
return self.score(pooled).squeeze(-1).masked_fill(~valid,-1e9)
|