Merge pull request #30 from kfu02/snd-qmix-ns

add snd for qmix non ps and indepedent agent env logic

Author: shalin-jain

baselines/QLearning/qmix.py

@@ -57,26 +57,26 @@ class MixingNetwork(nn.Module):
 
     @nn.compact
     def __call__(self, q_vals, states):
-        
+
         n_agents, time_steps, batch_size = q_vals.shape
         q_vals = jnp.transpose(q_vals, (1, 2, 0)) # (time_steps, batch_size, n_agents)
-        
+
         # hypernetwork
         w_1 = HyperNetwork(hidden_dim=self.hypernet_hidden_dim, output_dim=self.embedding_dim*n_agents, init_scale=self.init_scale)(states)
         b_1 = nn.Dense(self.embedding_dim, kernel_init=orthogonal(self.init_scale), bias_init=constant(0.))(states)
         w_2 = HyperNetwork(hidden_dim=self.hypernet_hidden_dim, output_dim=self.embedding_dim, init_scale=self.init_scale)(states)
         b_2 = HyperNetwork(hidden_dim=self.embedding_dim, output_dim=1, init_scale=self.init_scale)(states)
-        
+
         # monotonicity and reshaping
         w_1 = jnp.abs(w_1.reshape(time_steps, batch_size, n_agents, self.embedding_dim))
         b_1 = b_1.reshape(time_steps, batch_size, 1, self.embedding_dim)
         w_2 = jnp.abs(w_2.reshape(time_steps, batch_size, self.embedding_dim, 1))
         b_2 = b_2.reshape(time_steps, batch_size, 1, 1)
-    
+
         # mix
         hidden = nn.elu(jnp.matmul(q_vals[:, :, None, :], w_1) + b_1)
         q_tot  = jnp.matmul(hidden, w_2) + b_2
-        
+
         return q_tot.squeeze() # (time_steps, batch_size)
 
 
@@ -88,23 +88,23 @@ def __init__(self, start_e: float, end_e: float, duration: int):
         self.end_e    = end_e
         self.duration = duration
         self.slope    = (end_e - start_e) / duration
-        
+
     @partial(jax.jit, static_argnums=0)
     def get_epsilon(self, t: int):
         e = self.slope*t + self.start_e
         return jnp.clip(e, self.end_e)
-    
+
     @partial(jax.jit, static_argnums=0)
     def choose_actions(self, q_vals: dict, t: int, rng: chex.PRNGKey):
-        
+
         def explore(q, eps, key):
             key_a, key_e   = jax.random.split(key, 2) # a key for sampling random actions and one for picking
-            greedy_actions = jnp.argmax(q, axis=-1) # get the greedy actions 
+            greedy_actions = jnp.argmax(q, axis=-1) # get the greedy actions
             random_actions = jax.random.randint(key_a, shape=greedy_actions.shape, minval=0, maxval=q.shape[-1]) # sample random actions
             pick_random    = jax.random.uniform(key_e, greedy_actions.shape)<eps # pick which actions should be random
             chosed_actions = jnp.where(pick_random, random_actions, greedy_actions)
             return chosed_actions
-        
+
         eps = self.get_epsilon(t)
         keys = dict(zip(q_vals.keys(), jax.random.split(rng, len(q_vals)))) # get a key for each agent
         chosen_actions = jax.tree.map(lambda q, k: explore(q, eps, k), q_vals, keys)
@@ -128,7 +128,7 @@ def make_train(config, log_train_env, log_test_env, viz_test_env, env_name="MPE_
         config["TOTAL_TIMESTEPS"] // config["NUM_STEPS"] // config["NUM_ENVS"]
     )
 
-    
+
     def train(rng):
 
         # INIT ENV
@@ -166,7 +166,7 @@ def _env_sample_step(env_state, unused):
             sample_sequence_length=1,
             period=1,
         )
-        buffer_state = buffer.init(sample_traj_unbatched) 
+        buffer_state = buffer.init(sample_traj_unbatched)
 
         # INIT NETWORK
         # init agent
@@ -176,7 +176,7 @@ def _env_sample_step(env_state, unused):
             else:
                 exit("HyperMLP deprecated currently!") # TODO: to fix, pass in AGENT_HYPERNET_KWARGS
                 # agent = AgentHyperMLP(action_dim=wrapped_env.max_action_space, hidden_dim=config["AGENT_HIDDEN_DIM"], init_scale=config['AGENT_INIT_SCALE'], hypernet_hidden_dim=config["AGENT_HYPERNET_KWARGS"]["HIDDEN_DIM"], hypernet_init_scale=config["AGENT_HYPERNET_KWARGS"]["INIT_SCALE"], dim_capabilities=log_train_env.dim_capabilities)
-        else: 
+        else:
             if not config["AGENT_HYPERAWARE"]:
                 agent = AgentRNN(action_dim=wrapped_env.max_action_space, hidden_dim=config["AGENT_HIDDEN_DIM"], init_scale=config['AGENT_INIT_SCALE'])
             else:
@@ -290,7 +290,7 @@ def _env_step(step_state, unused):
                 dones_ = jax.tree.map(lambda x: x[np.newaxis, :], last_dones)
                 # get the q_values from the agent netwoek
                 hstate, q_vals = homogeneous_pass(params, hstate, obs_, dones_)
-                # remove the dummy time_step dimension and index qs by the valid actions of each agent 
+                # remove the dummy time_step dimension and index qs by the valid actions of each agent
                 valid_q_vals = jax.tree_util.tree_map(lambda q, valid_idx: q.squeeze(0)[..., valid_idx], q_vals, wrapped_env.valid_actions)
                 # explore with epsilon greedy_exploration
                 actions = explorer.choose_actions(valid_q_vals, t, key_a)
@@ -315,7 +315,7 @@ def _env_step(step_state, unused):
                 env_state,
                 init_obs,
                 init_dones,
-                hstate, 
+                hstate,
                 _rng,
                 time_state['timesteps'] # t is needed to compute epsilon
             )
@@ -360,12 +360,12 @@ def _loss_fn(params, target_network_params, init_hstate, learn_traj):
 
                 # compute q_tot with the mixer network
                 chosen_action_qvals_mix = mixer.apply(
-                    params['mixer'], 
+                    params['mixer'],
                     jnp.stack(list(chosen_action_qvals.values())),
                     learn_traj.obs['__all__'][:-1] # avoid last timestep
                 )
                 target_max_qvals_mix = mixer.apply(
-                    target_network_params['mixer'], 
+                    target_network_params['mixer'],
                     jnp.stack(list(target_max_qvals.values())),
                     learn_traj.obs['__all__'][1:] # avoid first timestep
                 )
@@ -399,7 +399,7 @@ def _td_lambda_target(ret, values):
                         + config['GAMMA']*(1-learn_traj.dones['__all__'][:-1])*target_max_qvals_mix
                     )
                     loss = jnp.mean((chosen_action_qvals_mix - jax.lax.stop_gradient(targets))**2)
-                
+
                 return loss
 
 
@@ -537,15 +537,15 @@ def _greedy_env_step(step_state, unused):
                 env_state,
                 init_obs,
                 init_dones,
-                hstate, 
+                hstate,
                 _rng,
             )
             step_state, (rewards, dones, infos, viz_env_states, obs, hstate) = jax.lax.scan(
                 _greedy_env_step, step_state, None, config["NUM_STEPS"]
             )
 
-            # get snd, NOTE: dim_c multiplier is currently hardcoded since it works for both fire and transport 
-            snd_value = snd(rollouts=obs, hiddens=hstate, dim_c=len(test_env.training_agents)*2, params=params, alg='qmix', agent=agent)
+            # get snd, NOTE: dim_c multiplier is currently hardcoded since it works for both fire and transport
+            snd_value = snd(rollouts=obs, hiddens=hstate, dim_c=len(test_env.training_agents)*2, params=params, alg='qmix' if config["PARAMETERS_SHARING"] else 'qmix_ns', agent=agent)
 
             def fire_env_metrics(final_env_state):
                 """
@@ -635,7 +635,7 @@ def callback(timestep, val):
                     print(f"Timestep: {timestep}, return: {val}")
                 jax.debug.callback(callback, time_state['timesteps']*config['NUM_ENVS'], first_returns['__all__'].mean())
             return {"metrics": metrics, "viz_env_states": viz_env_states}
-        
+
         time_state = {
             'timesteps':jnp.array(0),
             'updates':  jnp.array(0)
@@ -662,7 +662,7 @@ def callback(timestep, val):
             _update_step, runner_state, None, config["NUM_UPDATES"]
         )
         return {'runner_state':runner_state, 'metrics':metrics}
-    
+
     return train
 
 @hydra.main(version_base=None, config_path="./config", config_name="config")
@@ -673,7 +673,7 @@ def main(config):
 
     env_name = config["env"]["ENV_NAME"]
     alg_name = f'qmix_{"ps" if config["alg"].get("PARAMETERS_SHARING", True) else "ns"}'
-    
+
     # smac init neeeds a scenario
     if 'smax' in env_name.lower():
         config['env']['ENV_KWARGS']['scenario'] = map_name_to_scenario(config['env']['MAP_NAME'])
@@ -688,7 +688,7 @@ def main(config):
         log_test_env = LogWrapper(viz_test_env)
 
     config["alg"]["NUM_STEPS"] = config["alg"].get("NUM_STEPS", train_env.max_steps) # default steps defined by the env
-    
+
     hyper_tag = "hyper" if config["alg"]["AGENT_HYPERAWARE"] else "normal"
     recurrent_tag = "RNN" if config["alg"]["AGENT_RECURRENT"] else "MLP"
     aware_tag = "aware" if config["env"]["ENV_KWARGS"]["capability_aware"] else "unaware"
@@ -714,12 +714,12 @@ def main(config):
         config=config,
         mode=config["WANDB_MODE"],
     )
-    
+
     rng = jax.random.PRNGKey(config["SEED"])
     rngs = jax.random.split(rng, config["NUM_SEEDS"])
     train_vjit = jax.jit(jax.vmap(make_train(config["alg"], log_train_env, log_test_env, viz_test_env, env_name=config["env"]["ENV_NAME"])))
     outs = jax.block_until_ready(train_vjit(rngs))
-    
+
     # save params
     if config['SAVE_PATH'] is not None:
 
@@ -779,4 +779,4 @@ def save_params(params: Dict, filename: Union[str, os.PathLike]) -> None:
 
 if __name__ == "__main__":
     main()
-    
+

jaxmarl/environments/mpe/simple.py

@@ -1,4 +1,4 @@
-""" 
+"""
 Base class for MPE PettingZoo envs.
 
 TODO: viz for communication env, e.g. crypto
@@ -52,6 +52,7 @@ def __init__(
     ):
         self.test_env_flag = kwargs["test_env_flag"] if "test_env_flag" in kwargs else False
         self.test_capabilities = kwargs["test_capabilities"] if "test_capabilities" in kwargs else None
+        self.independent_agents = kwargs["independent_agents"] if "independent_agents" in kwargs else None
 
         # Agent and entity constants
         self.num_agents = num_agents
@@ -136,7 +137,7 @@ def __init__(
             self.agent_accels = kwargs["agent_accels"]
             # assert (len(self.agent_accels) >= self.num_agents), f"Not enough agent_accels, {len(self.agent_accels)} < {self.num_agents}"
             self.agent_accels = jnp.array(self.agent_accels)
-        
+
         if "agent_capacities" in kwargs:
             self.agent_capacities = kwargs["agent_capacities"]
             self.agent_capacities = jnp.array(self.agent_capacities)
@@ -297,7 +298,7 @@ def reset(self, key: chex.PRNGKey) -> Tuple[chex.Array, State]:
         # if self.test_env_flag and self.test_capabilities is not None:
         #     team_capabilities = jnp.asarray(self.test_capabilities)
 
-        
+
         agent_rads = self.agent_rads[selected_agents]
         agent_accels = self.agent_accels[selected_agents]
         agent_capacities = self.agent_capacities[selected_agents] if self.agent_capacities else np.zeros((self.num_agents, 2))
@@ -521,7 +522,7 @@ def map_bounds_reward(self, x: float):
         m = x < 1.0
         mr = (x - 0.9) * 10
         br = jnp.min(jnp.array([jnp.exp(2 * x - 2), 10]))
-        return jax.lax.select(m, mr, br) * ~w   
+        return jax.lax.select(m, mr, br) * ~w
 
 
 if __name__ == "__main__":

jaxmarl/environments/mpe/simple_fire.py

@@ -30,7 +30,7 @@ def __init__(
         fire_pos_dim = num_landmarks * 2
         fire_rad_dim = num_landmarks
         observation_spaces = {
-            i:Box(-jnp.inf, jnp.inf, (pos_dim + vel_dim + self.dim_capabilities + fire_pos_dim + fire_rad_dim)) 
+            i:Box(-jnp.inf, jnp.inf, (pos_dim + vel_dim + self.dim_capabilities + fire_pos_dim + fire_rad_dim))
             for i in agents
         }
 
@@ -40,7 +40,7 @@ def __init__(
         # env specific parameters
         self.test_teams = jnp.array(kwargs["test_teams"]) if "test_teams" in kwargs else None
         self.fire_rad_range = kwargs["fire_rad_range"] if "fire_rad_range" in kwargs else [0.2, 0.3]
-        
+
         # reward shaping
         self.fire_out_reward = kwargs["fire_out_reward"] if "fire_out_reward" in kwargs else 1
         self.uncovered_penalty_factor = kwargs["uncovered_penalty_factor"] if "uncovered_penalty_factor" in kwargs else 2
@@ -207,7 +207,7 @@ def _spawn_one_fire(carry, _):
 
             # if new fire spawn is valid, add it to the fire list, and incr the fire index
             new_fires = jax.lax.cond(
-                new_fire_valid, 
+                new_fire_valid,
                 lambda: new_fire_added, # T
                 lambda: existing_fires, # F
             )
@@ -240,18 +240,23 @@ def _spawn_one_fire(carry, _):
             ]
         )
 
-        # randomly sample N_agents' capabilities from the possible agent pool (hence w/out replacement)
-        selected_agents = jax.random.choice(key_c, self.num_agents, shape=(self.num_agents,), replace=False)
-        agent_rads = self.agent_rads[selected_agents]
-        agent_accels = self.agent_accels[selected_agents]
-
-        # unless a test distribution is provided and this is a test_env
-        if self.test_env_flag and self.test_teams is not None:
-            # pick one of the test teams at random
-            selected_team = jax.random.choice(key_tt, self.test_teams.shape[0], shape=(1,))
-            test_team = self.test_teams[selected_team].squeeze()
-            agent_rads = test_team[0::2]
-            agent_accels = test_team[1::2]
+        if self.independent_agents:
+            # if independent policies do not sample teams and capabilities, keep constant
+            # NOTE: assumes that agent_rad and agent_accels are n_agent length
+            agent_rads = self.agent_rads
+            agent_accels = self.agent_accels
+        else:
+            # randomly sample N_agents' capabilities from the possible agent pool (hence w/out replacement)
+            selected_agents = jax.random.choice(key_c, self.num_agents, shape=(self.num_agents,), replace=False)
+            agent_rads = self.agent_rads[selected_agents]
+            agent_accels = self.agent_accels[selected_agents]
+            # unless a test distribution is provided and this is a test_env
+            if self.test_env_flag and self.test_teams is not None:
+                # pick one of the test teams at random
+                selected_team = jax.random.choice(key_tt, self.test_teams.shape[0], shape=(1,))
+                test_team = self.test_teams[selected_team].squeeze()
+                agent_rads = test_team[0::2]
+                agent_accels = test_team[1::2]
 
         state = State(
             p_pos=p_pos,