Views
No views yet
1import gymnasium as gym
2import os
3from tqdm import tqdm
4import torch
5
6
7
8class CartPole(torch.nn.Module):
9 def __init__(self,):
10 super(CartPole, self).__init__()
11 self.model = torch.nn.Sequential(
12 torch.nn.Linear(4,64),
13 torch.nn.ReLU(),
14 torch.nn.Linear(64,2),
15 )
16
17 def forward(self, x):
18 x = self.model(x)
19 return x
20
21
22def run(model, episodes):
23 video_length = episodes
24 env = gym.make("CartPole-v1", render_mode="human") # human, rgb_array
25 obs, _ = env.reset()
26 total_reward = 0.0
27 with torch.no_grad():
28 for i in tqdm(range(video_length+1)):
29 x = torch.tensor(obs).float().unsqueeze(0).to('cuda')
30 action = model(x).argmax(dim=-1).item()
31
32 obs, reward, terminated, truncated, info = env.step(action)
33
34 if terminated or truncated:
35 obs, _ = env.reset()
36 total_reward+=reward
37
38 env.close()
39 print(f"total reward : {total_reward}")
40
41
42
43
44model = CartPole()
45model.load_state_dict(torch.load(
46 os.path.join(os.getcwd(),"99.61_99_policy_net.pth")
47)['model_state_dict'])
48model.to("cuda")
49model.eval()
50run(model=model, episodes=500)