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)
|