Views
No views yet
1import torch
2import json
3import gymnasium as gym
4
5from agent import SimpleAgent
6from environment import make_env
7
8#Load the configuration
9with open("config.json", "r") as f:
10 config = json.load(f)
11
12env_id = config["env_id"]
13hidden_dim = config["hidden_dim"]
14
15# Create environment. Get action and space dimensions
16env, state_size, action_size = make_env(
17 env_id,
18 render_mode="human",
19 normalise_obs=config["normalise_obs"],
20)
21
22# Instantiate the agent and load the trained policy network
23agent = SimpleAgent(state_size, action_size, hidden_dim)
24agent.policy.load_state_dict(torch.load("model.pt"))
25
26# Enjoy the agent!
27state, _ = env.reset()
28done = False
29
30while not done:
31 action = agent.select_action(state)
32 state, reward, terminated, truncated, _ = env.step(action)
33
34 done = terminated or truncated
35
36 env.render()
37
38env.close()