feat: add continuous action space support

comrad/envs/action_space.py

@@ -328,6 +328,40 @@ def doom_action_space_stag_hunt():
         )
     )
 
+def doom_action_space_stag_hunt_continuous():
+    """
+    MOVE_FORWARD
+    MOVE_BACKWARD
+    MOVE_RIGHT
+    MOVE_LEFT
+    ATTACK
+    TURN_LEFT_RIGHT_DELTA
+    """
+    return gym.spaces.Tuple(
+        (
+            Discrete(3),  # noop, forward, backward
+            Discrete(3),  # noop, move right, move left
+            Discrete(2),  # noop, attack
+            Box(np.float32(-1.0), np.float32(1.0), (1,)), # continuous turning
+        )
+    )
+
+def doom_action_space_stag_hunt_continuous_full():
+    """
+    MOVE_FORWARD_BACKWARD_DELTA
+    MOVE_LEFT_RIGHT_DELTA
+    ATTACK
+    TURN_LEFT_RIGHT_DELTA
+    """
+    return gym.spaces.Tuple(
+        (
+            Box(np.float32(-1.0), np.float32(1.0), (1,)),  # continuous movement fb
+            Box(np.float32(-1.0), np.float32(1.0), (1,)),  # continuous movement lr
+            Discrete(2),  # noop, attack
+            Box(np.float32(-1.0), np.float32(1.0), (1,)),  # continuous turning
+        )
+    )
+
 def doom_action_space_lava_maze():
     """
     MOVE_FORWARD

comrad/envs/doom_gym.py

@@ -429,27 +429,37 @@ def _convert_actions(self, actions):
             actions = (actions,)
 
         actions_flattened = []
-        for i, action in enumerate(actions):
-            if isinstance(spaces[i], Discretized):
-                # discretized continuous action
-                # check discretized first because it's a subclass of gym.spaces.Discrete
-                # the order of if clauses here matters! DON'T CHANGE THE ORDER OF IFS!
-
-                continuous_action = spaces[i].to_continuous(action)
-                actions_flattened.append(continuous_action)
-            elif isinstance(spaces[i], gym.spaces.Discrete):
-                # standard discrete action
-                num_non_idle_actions = spaces[i].n - 1
-                action_one_hot = np.zeros(num_non_idle_actions, dtype=np.uint8)
-                if action > 0:
-                    action_one_hot[action - 1] = 1  # 0th action in each subspace is a no-op
-
-                actions_flattened.extend(action_one_hot)
-            elif isinstance(spaces[i], gym.spaces.Box):
-                # continuous action
-                actions_flattened.extend(list(action * self.delta_actions_scaling_factor))
-            else:
-                raise NotImplementedError(f"Action subspace type {type(spaces[i])} is not supported!")
+        try:
+            for i, action in enumerate(actions):
+                if isinstance(spaces[i], Discretized):
+                    # discretized continuous action
+                    # check discretized first because it's a subclass of gym.spaces.Discrete
+                    # the order of if clauses here matters! DON'T CHANGE THE ORDER OF IFS!
+
+                    continuous_action = spaces[i].to_continuous(action)
+                    actions_flattened.append(continuous_action)
+                elif isinstance(spaces[i], gym.spaces.Discrete):
+                    # standard discrete action
+                    num_non_idle_actions = spaces[i].n - 1
+                    action_one_hot = np.zeros(num_non_idle_actions, dtype=np.uint8)
+                    if action > 0:
+                        action_one_hot[action - 1] = 1  # 0th action in each subspace is a no-op
+
+                    actions_flattened.extend(action_one_hot)
+                elif isinstance(spaces[i], gym.spaces.Box):
+                    # continuous action
+                    val = action * self.delta_actions_scaling_factor
+                    if isinstance(val, (int, float, np.number)):
+                        actions_flattened.append(float(val))
+                    else:
+                        actions_flattened.extend(list(val))
+                else:
+                    raise NotImplementedError(f"Action subspace type {type(spaces[i])} is not supported!")
+        except IndexError as e:
+            log.error(f"IndexError in _convert_actions! len(spaces)={len(spaces)}, len(actions)={len(actions)}")
+            log.error(f"actions: {actions}")
+            log.error(f"spaces: {spaces}")
+            raise e
 
         return actions_flattened
 

comrad/scenarios/stag_hunt_arena_continuous.cfg

@@ -0,0 +1,38 @@
+doom_scenario_path = stag_hunt_arena.wad
+
+living_reward = 0
+
+screen_resolution = RES_320X240
+screen_format = CRCGCB
+render_hud = true
+render_crosshair = true
+render_weapon = true
+render_decals = false
+render_particles = false
+window_visible = false
+
+# 4500 tics ~= 2.1 minutes
+episode_timeout = 4500
+
+available_buttons =
+    {
+        MOVE_FORWARD
+        MOVE_BACKWARD
+        MOVE_RIGHT
+        MOVE_LEFT
+        ATTACK
+        TURN_LEFT_RIGHT_DELTA
+    }
+
+# Game variables that will be in the state
+# USER51: sh_stag_health
+# USER52: sh_stag_vulnerable
+# USER53: sh_rabbit_kills
+# USER54: sh_stag_kills
+# USER55: sh_stag_alive
+# USER56: sh_p1_rabbit_kills
+# USER57: sh_p2_rabbit_kills
+available_game_variables = { HEALTH KILLCOUNT POSITION_X POSITION_Y USER51 USER52 USER53 USER54 USER55 USER56 USER57 }
+
+mode = PLAYER
+doom_skill = 3

comrad/scenarios/stag_hunt_arena_continuous_full.cfg

@@ -0,0 +1,36 @@
+doom_scenario_path = stag_hunt_arena.wad
+
+living_reward = 0
+
+screen_resolution = RES_320X240
+screen_format = CRCGCB
+render_hud = true
+render_crosshair = true
+render_weapon = true
+render_decals = false
+render_particles = false
+window_visible = false
+
+# 4500 tics ~= 2.1 minutes
+episode_timeout = 4500
+
+available_buttons =
+    {
+        MOVE_FORWARD_BACKWARD_DELTA
+        MOVE_LEFT_RIGHT_DELTA
+        ATTACK
+        TURN_LEFT_RIGHT_DELTA
+    }
+
+# Game variables that will be in the state
+# USER51: sh_stag_health
+# USER52: sh_stag_vulnerable
+# USER53: sh_rabbit_kills
+# USER54: sh_stag_kills
+# USER55: sh_stag_alive
+# USER56: sh_p1_rabbit_kills
+# USER57: sh_p2_rabbit_kills
+available_game_variables = { HEALTH KILLCOUNT POSITION_X POSITION_Y USER51 USER52 USER53 USER54 USER55 USER56 USER57 }
+
+mode = PLAYER
+doom_skill = 3

comrad/utils/doom_utils.py

@@ -27,6 +27,8 @@
     doom_action_space_dumb_enemies,
     doom_action_space_stealth_labyrinth,
     doom_action_space_stag_hunt,
+    doom_action_space_stag_hunt_continuous,
+    doom_action_space_stag_hunt_continuous_full,
     doom_action_space_coop_health_gathering,
     doom_action_space_foraging_commons,
     doom_action_space_rhythm_sync,
@@ -251,6 +253,30 @@ def __init__(
         shared_reward_alpha=0.0,
     ),
 
+    DoomSpec(
+        "stag_hunt_arena_continuous",
+        "stag_hunt_arena_continuous.cfg",
+        doom_action_space_stag_hunt_continuous(),
+        1.0,
+        4500,
+        num_agents=2,
+        forcerespawn=1,
+        extra_wrappers=[(StagHuntArenaRewardShaping, {})],
+        shared_reward_alpha=0.0,
+    ),
+
+    DoomSpec(
+        "stag_hunt_arena_continuous_full",
+        "stag_hunt_arena_continuous_full.cfg",
+        doom_action_space_stag_hunt_continuous_full(),
+        1.0,
+        4500,
+        num_agents=2,
+        forcerespawn=1,
+        extra_wrappers=[(StagHuntArenaRewardShaping, {})],
+        shared_reward_alpha=0.0,
+    ),
+
     DoomSpec(
         "ammo_carrier",
         "ammo_carrier.cfg",

pyproject.toml

@@ -7,7 +7,7 @@ name = "comrad-bench"
 version = "1.0.0"
 description = "MARL benchmark suite in ViZDoom"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.11,<3.13"
 license = {text = "MIT"}
 authors = [
     {name = "Khoi H.B. Nguyen", email = "baokhoi136@gmail.com"},
@@ -54,7 +54,7 @@ dependencies = [
     "seaborn>=0.13.2",
     "scipy>=1.17.0",
     "vizdoom>=1.3.0",
-    "ml-dtypes>=0.4.0,<0.5",
+    "ml-dtypes>=0.5.0",
 ]
 
 [project.urls]

sample_factory/algo/sampling/batched_sampling.py

@@ -54,6 +54,10 @@ def preprocess_actions(
             )
         # this line can be used to transpose the actions, perhaps add as an option ?
         # out_actions = list(zip(*out_actions)) # transpose
+        # Yeah it can be used for cont. action vect
+        if len(out_actions) > 0 and isinstance(out_actions[0], (np.ndarray, Tensor)):
+            out_actions = list(zip(*out_actions))
+
         return out_actions
 
     raise NotImplementedError(f"Unknown action space type: {env_info.action_space}")
@@ -403,7 +407,7 @@ def advance_rollouts(self, policy_id: PolicyID, timing) -> Tuple[List[Dict], Lis
 
         self.env_step_ready = True
         return complete_rollouts, episodic_stats
-    
+
     def _get_task_idx_buffer(self, infos) -> torch.Tensor:
         """Extract task_idx from infos into a [num_agents] int32 tensor."""
         buf = torch.full(

sample_factory/algo/utils/action_distributions.py

@@ -336,6 +336,10 @@ def kl_divergence(self, other):
         kl = torch.distributions.kl.kl_divergence(self, other)
         return kl
 
+    def symmetric_kl_with_uniform_prior(self):
+        kl = 0.5 * (self.means.pow(2) + (2 * self.log_std).exp() - 2 * self.log_std - 1)
+        return kl.sum(dim=-1)
+
     def summaries(self):
         return dict(
             action_mean=self.means.mean(),

sample_factory/tests/envs/tuple_action_envs/test_mixed_dist_env_batched.py

@@ -74,11 +74,16 @@ def reset(self, **kwargs):
             infos.append(info)
         return obss, infos
 
-    def step(self, action: List[np.ndarray]):
+    def step(self, action):
         obss, rewards, terms, truncs, infos = [], [], [], [], []
 
+        subspace_major = len(action) > 0 and isinstance(action[0], np.ndarray)
         for i, env in enumerate(self.envs):
-            obs, reward, terminated, truncated, info = env.step([action[0][i], action[1][i]])
+            if subspace_major:
+                actions_i = [action[j][i] for j in range(len(action))]
+            else:
+                actions_i = list(action[i])
+            obs, reward, terminated, truncated, info = env.step(actions_i)
             obss.append(obs),
             rewards.append(reward)
             terms.append(terminated)