File size: 663 Bytes
18ee2df
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
# pip install ale-py gymnasium torch safetensors huggingface_hub
# You also need the phoenix_rl package that trained this model (it defines the network class).
import json, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from phoenix_rl.config import Cfg
from phoenix_rl.nets import ActorCritic
repo = "<your-username>/ppo-phoenix"
cfg = Cfg(**json.load(open(hf_hub_download(repo, "config.json"))))
net = ActorCritic(8, cfg.arch, cfg.hidden); net.load_state_dict(load_file(hf_hub_download(repo, "model.safetensors"))); net.eval()
# obs: uint8 tensor (B, 4, 84, 84);  logits, value = net(obs);  action = logits.argmax(-1)