Fix-configs

mighty/configs/environment/pufferlib_ocean/password.yaml

@@ -4,4 +4,4 @@ num_steps: 50_000
 env: pufferlib.ocean.password
 env_kwargs: {}
 env_wrappers: []
-num_envs: 1
+num_envs: 64

mighty/configs/environment/pufferlib_ocean/squared.yaml

@@ -3,5 +3,5 @@
 num_steps: 50_000 
 env: pufferlib.ocean.squared
 env_kwargs: {}
-env_wrappers: [mighty.utils.wrappers.FlattenVecObs]
-num_envs: 1
+env_wrappers: [mighty.mighty_utils.wrappers.FlattenVecObs]
+num_envs: 64

mighty/configs/environment/pufferlib_ocean/stochastic.yaml

@@ -4,4 +4,4 @@ num_steps: 50_000
 env: pufferlib.ocean.stochastic
 env_kwargs: {}
 env_wrappers: []
-num_envs: 1
+num_envs: 64

mighty/configs/exploration/ez_greedy.yaml

@@ -1,3 +1,4 @@
 # @package _global_
 algorithm_kwargs:
-  policy_class: mighty.mighty_exploration.EZGreedy
+  policy_class: mighty.mighty_exploration.EZGreedy
+  policy_kwargs: null

mighty/mighty_agents/base_agent.py

@@ -13,7 +13,7 @@
 import pandas as pd
 import torch
 import wandb
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
 from rich import print
 from rich.layout import Layout
 from rich.live import Live
@@ -323,6 +323,10 @@ def initialize_agent(self) -> None:
         if isinstance(self.buffer_class, type) and issubclass(
             self.buffer_class, PrioritizedReplay
         ):
+            if isinstance(self.buffer_kwargs, DictConfig):
+                self.buffer_kwargs = OmegaConf.to_container(
+                    self.buffer_kwargs, resolve=True
+                )
             # 1) Get observation-space shape
             try:
                 obs_space = self.env.single_observation_space

mighty/mighty_utils/wrappers.py

@@ -106,19 +106,21 @@ def __init__(self, env):
 
         """
         super().__init__(env)
-        self.n_actions = len(self.env.single_action_space.nvec)
-        self.single_action_space = gym.spaces.Discrete(
-            np.prod(self.env.single_action_space.nvec)
-        )
+        self.n_actions = len(self.env.action_space.nvec)
+
         self.action_mapper = {}
         for idx, prod_idx in zip(
-            range(np.prod(self.env.single_action_space.nvec)),
+            range(np.prod(self.env.action_space.nvec)),
             itertools.product(
-                *[np.arange(val) for val in self.env.single_action_space.nvec]
+                *[np.arange(val) for val in self.env.action_space.nvec]
             ),
         ):
             self.action_mapper[idx] = prod_idx
 
+        self.action_space = gym.spaces.Discrete(
+            int(np.prod(self.env.action_space.nvec))
+        )
+
     def step(self, action):
         """Maps discrete action value to array."""
         action = [self.action_mapper[a] for a in action]