Views
No views yet
1import torch
2import gymnasium as gym
3import torch.nn as nn
4import torch.nn.functional as F
5from torch.distributions import Categorical
6import numpy as np
7
8# Define the Actor network
9class Actor(nn.Module):
10 def __init__(self, state_dim, action_dim, hidden_size=64):
11 super().__init__()
12 self.network = nn.Sequential(
13 nn.Linear(state_dim, hidden_size),
14 nn.Tanh(),
15 nn.Linear(hidden_size, hidden_size),
16 nn.Tanh(),
17 nn.Linear(hidden_size, action_dim)
18 )
19
20 def forward(self, x):
21 return self.network(x)
22
23# Load the model
24checkpoint = torch.load("model.pt", map_location='cpu')
25actor = Actor(state_dim=8, action_dim=4, hidden_size=checkpoint['config']['hidden_size'])
26actor.load_state_dict(checkpoint['actor_state_dict'])
27actor.eval()
28
29# Test the agent
30env = gym.make("LunarLander-v2")
31state, _ = env.reset()
32total_reward = 0
33
34for _ in range(1000): # Max steps
35 with torch.no_grad():
36 state_tensor = torch.FloatTensor(state).unsqueeze(0)
37 logits = actor(state_tensor)
38 action = torch.argmax(logits, dim=-1).item()
39
40 state, reward, terminated, truncated, _ = env.step(action)
41 total_reward += reward
42
43 if terminated or truncated:
44 break
45
46print(f"Total reward: {total_reward:.2f}")