Views
No views yet
1def pong_obs_modification(obs, _space, player_id):
2 obs[:9, :, :] = 0
3 if "second" in player_id:
4 # Mirror the image
5 obs = obs[:, ::-1, :]
6 return obs
7
8
9def get_env(args, run_name):
10 env = importlib.import_module(f"pettingzoo.atari.{args.env_id}").parallel_env()
11 env = ss.max_observation_v0(env, 2)
12 env = ss.frame_skip_v0(env, 4)
13 env = ss.clip_reward_v0(env, lower_bound=-1, upper_bound=1)
14 env = ss.color_reduction_v0(env, mode="B")
15 env = ss.resize_v1(env, x_size=84, y_size=84)
16 env = ss.frame_stack_v1(env, 4)
17 # Remove the score from the observation
18 if "pong" in args.env_id:
19 env = ss.lambda_wrappers.observation_lambda_v0(
20 env,
21 pong_obs_modification,
22 )
23 # env = ss.agent_indicator_v0(env, type_only=False)
24 env = ss.pettingzoo_env_to_vec_env_v1(env)
25 envs = ss.concat_vec_envs_v1(env, args.num_envs // 2, num_cpus=0, base_class="gym")
26 envs.single_observation_space = envs.observation_space
27 envs.single_action_space = envs.action_space
28 envs.is_vector_env = True
29 envs = gym.wrappers.RecordEpisodeStatistics(envs)
30 if args.capture_video:
31 envs = gym.wrappers.RecordVideo(envs, f"videos/{run_name}")
32 assert isinstance(
33 envs.single_action_space, gym.spaces.Discrete
34 ), "only discrete action space is supported"
35 return envs1def atari_network(orth_init=False):
2 init = layer_init if orth_init else lambda m: m
3 return nn.Sequential(
4 init(nn.Conv2d(4, 32, 8, stride=4)),
5 nn.ReLU(),
6 init(nn.Conv2d(32, 64, 4, stride=2)),
7 nn.ReLU(),
8 init(nn.Conv2d(64, 64, 3, stride=1)),
9 nn.ReLU(),
10 nn.Flatten(),
11 init(nn.Linear(64 * 7 * 7, 512)),
12 nn.ReLU(),
13 )
14
15class Agent(nn.Module):
16 def __init__(self, envs, share_network=False):
17 super().__init__()
18 self.actor_network = atari_network(orth_init=True)
19 self.share_network = share_network
20 if share_network:
21 self.critic_network = self.actor_network
22 else:
23 self.critic_network = atari_network(orth_init=True)
24 self.actor = layer_init(nn.Linear(512, envs.single_action_space.n), std=0.01)
25 self.critic = layer_init(nn.Linear(512, 1), std=1)
26
27 def get_value(self, x):
28 x = x.clone()
29 x[:, :, :, [0, 1, 2, 3]] /= 255.0
30 return self.critic(self.critic_network(x.permute((0, 3, 1, 2))))
31
32 def get_action_and_value(self, x, action=None):
33 x = x.clone()
34 x[:, :, :, [0, 1, 2, 3]] /= 255.0
35 logits = self.actor(self.actor_network(x.permute((0, 3, 1, 2))))
36 probs = Categorical(logits=logits)
37 if action is None:
38 action = probs.sample()
39 return (
40 action,
41 probs.log_prob(action),
42 probs.entropy(),
43 self.critic(self.critic_network(x.permute((0, 3, 1, 2)))),
44 )
45
46 def load(self, path):
47 self.load_state_dict(torch.load(path))
48 if self.share_network:
49 self.critic_network = self.actor_network