[IBR-2068] Add ppo with discrete action

Author: isk03276

configs/lunarlander_v2/ppo.yaml

@@ -0,0 +1,40 @@
+type: "PPOAgent"
+hyper_params:
+  gamma: 0.99
+  tau: 0.95
+  batch_size: 32
+  max_epsilon: 0.2
+  min_epsilon: 0.2
+  epsilon_decay_period: 1500
+  w_value: 1.0
+  w_entropy: 0.001
+  gradient_clip_ac: 0.5
+  gradient_clip_cr: 1.0
+  epoch: 16
+  rollout_len: 256
+  n_workers: 12
+  use_clipped_value_loss: False
+  standardize_advantage: True
+
+learner_cfg:
+  type: "PPOLearner"
+  backbone:
+    actor:
+    critic:
+    shared_actor_critic:
+  head:
+    actor:
+      type: "CategoricalDist"
+      configs: 
+        hidden_sizes: [256, 256]
+        output_activation: "identity"
+    critic:
+      type: "MLP"
+      configs:
+        hidden_sizes: [256, 256]
+        output_size: 1
+        output_activation: "identity"
+  optim_cfg:
+    lr_actor: 0.0003
+    lr_critic: 0.001
+    weight_decay: 0.0

configs/pong_no_frameskip_v4/ppo.yaml

@@ -0,0 +1,47 @@
+type: "PPOAgent"
+hyper_params:
+  gamma: 0.99
+  tau: 0.95
+  batch_size: 32
+  max_epsilon: 0.2
+  min_epsilon: 0.2
+  epsilon_decay_period: 1500
+  w_value: 1.0
+  w_entropy: 0.001
+  gradient_clip_ac: 0.5
+  gradient_clip_cr: 1.0
+  epoch: 16
+  rollout_len: 256
+  n_workers: 4
+  use_clipped_value_loss: False
+  standardize_advantage: True
+
+learner_cfg:
+  type: "PPOLearner"
+  backbone:
+    actor:
+    critic:
+    shared_actor_critic:
+      type: "CNN"
+      configs:
+        input_sizes: [4, 32, 64]
+        output_sizes: [32, 64, 64]
+        kernel_sizes: [8, 4, 3]
+        strides: [4, 2, 1]
+        paddings: [1, 0, 0]
+  head:
+    actor:
+      type: "CategoricalDist"
+      configs: 
+        hidden_sizes: [512]
+        output_activation: "identity"
+    critic:
+      type: "MLP"
+      configs:
+        hidden_sizes: [512]
+        output_size: 1
+        output_activation: "identity"
+  optim_cfg:
+    lr_actor: 0.0003
+    lr_critic: 0.001
+    weight_decay: 0.0

rl_algorithms/common/networks/heads.py

@@ -221,7 +221,6 @@ def __init__(
     def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
         """Forward method implementation."""
         ac_logits = super().forward(x)
-        # ac_probs = F.softmax(ac_logits, dim=-1)
 
         # get categorical distribution and action
         dist = Categorical(logits=ac_logits)

rl_algorithms/ppo/agent.py

@@ -123,7 +123,10 @@ def __init__(
         self.learner = build_learner(self.learner_cfg, build_args)
 
     def make_parallel_env(self, max_episode_steps, n_workers):
-        env_gen = env_generator(self.env.spec.id, max_episode_steps)
+        if "env_generator" in self.env_info.keys():
+            env_gen = self.env_info.env_generator
+        else:
+            env_gen = env_generator(self.env.spec.id, max_episode_steps)
         env_multi = make_envs(env_gen, n_envs=n_workers)
         return env_multi
 
@@ -135,7 +138,9 @@ def select_action(self, state: np.ndarray) -> torch.Tensor:
             log_prob = dist.log_prob(selected_action)
 
             if self.is_test:
-                selected_action = dist.mean
+                selected_action = (
+                    dist.logits.argmax() if self.is_discrete else dist.mean
+                )
 
             else:
                 _selected_action = (

run_pong_no_frameskip_v4.py

@@ -84,15 +84,25 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def env_generator(env_name, max_episode_steps, frame_stack):
+    def _thunk(rank: int):
+        env = atari_env_generator(env_name, max_episode_steps, frame_stack=frame_stack)
+        env.seed(777 + rank + 1)
+        return env
+
+    return _thunk
+
+
 def main():
     """Main."""
     args = parse_args()
 
     # env initialization
     env_name = "PongNoFrameskip-v4"
-    env = atari_env_generator(
+    env_gen = env_generator(
         env_name, args.max_episode_steps, frame_stack=args.framestack
     )
+    env = env_gen(0)
 
     # set a random seed
     common_utils.set_random_seed(args.seed, env)
@@ -112,6 +122,7 @@ def main():
         observation_space=env.observation_space,
         action_space=env.action_space,
         is_atari=True,
+        env_generator=env_gen,
     )
     log_cfg = dict(agent=cfg.agent.type, curr_time=curr_time, cfg_path=args.cfg_path)
     build_args = dict(