Skip to content

Commit b6734b6

Browse files
iamlilAJiamlilAJ
committed
clean
1 parent 7416038 commit b6734b6

8 files changed

Lines changed: 27 additions & 269 deletions

File tree

‎baselines/IPPO/ippo_pre.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -411,11 +411,12 @@ def wrapped_make_train():
411411
},
412412

413413
'parameters': {
414-
# 'TAU': {'values': [0.01, 0.1, 1.0]},
415-
"SEED": {
416-
"values": [50, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
417-
21, 22, 23, 24, 25, 26, 27, 28, 29, 30]
418-
}
414+
415+
"LR": {
416+
"values": [0.0005, 0.0006, 0.0007]
417+
418+
}
419+
419420
}
420421
}
421422

‎baselines/MAPPO/mappo_pre.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -482,8 +482,8 @@ def wrapped_make_train():
482482

483483
'parameters': {
484484

485-
"SEED": {
486-
"values": [30, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]
485+
"LR": {
486+
"values": [0.0005, 0.0006, 0.0007]
487487

488488
}
489489
}

‎baselines/QLearning/agent/iql_agent.py‎

Lines changed: 0 additions & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -57,28 +57,6 @@ def __call__(self, hidden, obs, dones):
5757

5858
return hidden, embedding
5959

60-
class RewardRNN(nn.Module):
61-
# homogenous agent for parameters sharing, assumes all agents have same obs and action dim
62-
63-
hidden_dim: int
64-
init_scale: float
65-
66-
@nn.compact
67-
def __call__(self, hidden, obs, dones):
68-
69-
embedding = nn.Dense(self.hidden_dim, kernel_init=orthogonal(self.init_scale), bias_init=constant(0.1))(obs)
70-
embedding = nn.relu(embedding)
71-
72-
rnn_in = (embedding, dones)
73-
hidden, embedding = ScannedRNN()(hidden, rnn_in)
74-
# output reward
75-
rew = nn.Dense(1, kernel_init=orthogonal(self.init_scale), bias_init=constant(0.1))(embedding)
76-
77-
rew = nn.tanh(rew)
78-
79-
return hidden, rew
80-
81-
8260
class IQLAgent(nn.Module):
8361
action_dim: int
8462
hidden_dim: int
@@ -188,149 +166,4 @@ def __call__(self, agent_hidden, obs, dones, _):
188166

189167

190168

191-
def test_iql_agent():
192-
"""
193-
Test the IQLAgent with actual implementations and verify output shapes and masking.
194-
"""
195-
# Define test parameters
196-
num_agents = 4
197-
timesteps = 5
198-
batch_size = 2
199-
obs_dim = 14 # Updated from 20 to 14
200-
action_dim = 3
201-
hidden_dim = 16
202-
num_proxy_agents = 2
203-
pre_policy_output_dim = 2
204-
pre_policy_hidden_dim = 12
205-
gnn_embedding_dim = 14
206-
init_scale = 1.0
207-
208-
config = {'GNN_EMBEDDING_DIM': gnn_embedding_dim}
209-
210-
# Initialize the agent
211-
agent = IQLAgent(
212-
action_dim=action_dim,
213-
hidden_dim=hidden_dim,
214-
num_agents=num_agents,
215-
num_proxy_agents=num_proxy_agents,
216-
init_scale=init_scale,
217-
pre_policy_output_dim=pre_policy_output_dim,
218-
pre_policy_hidden_dim=pre_policy_hidden_dim,
219-
config=config
220-
)
221-
222-
# Initialize parameters with a random key
223-
rng = jax.random.PRNGKey(0)
224-
225-
# Define the input shapes
226-
# agent_hidden: [num_agents, batch_size, hidden_dim]
227-
agent_hidden = jnp.zeros((num_agents, batch_size, hidden_dim))
228-
# obs: [num_agents, timesteps, batch_size, obs_dim=14]
229-
obs = jnp.ones((num_agents, timesteps, batch_size, obs_dim))
230-
# dones: [num_agents, timesteps, batch_size]
231-
dones = jnp.zeros((num_agents, timesteps, batch_size))
232-
# pre_policy_hidden: [num_agents, batch_size, pre_policy_hidden_dim]
233-
pre_policy_hidden = jnp.zeros((num_agents, batch_size, pre_policy_hidden_dim))
234-
235-
# Initialize the parameters using dummy inputs
236-
variables = agent.init(rng, agent_hidden, obs, dones, pre_policy_hidden)
237-
# Apply the agent
238-
agent_hidden_out, q_vals, pre_policy_hidden_out = agent.apply(
239-
variables, agent_hidden, obs, dones, pre_policy_hidden
240-
)
241-
242-
# Assertions to verify output shapes
243-
assert agent_hidden_out.shape == (num_agents, batch_size, hidden_dim), (
244-
f"agent_hidden_out shape mismatch: expected {(num_agents, batch_size, hidden_dim)}, "
245-
f"got {agent_hidden_out.shape}"
246-
)
247-
# assert q_vals.shape == (num_agents, batch_size, action_dim), (
248-
# f"q_vals shape mismatch: expected {(num_agents, batch_size, action_dim)}, "
249-
# f"got {q_vals.shape}"
250-
# )
251-
assert pre_policy_hidden_out.shape == (num_agents, batch_size, pre_policy_hidden_dim), (
252-
f"pre_policy_hidden_out shape mismatch: expected {(num_agents, batch_size, pre_policy_hidden_dim)}, "
253-
f"got {pre_policy_hidden_out.shape}"
254-
)
255-
256-
# To verify the masking, create a subclass that returns pre_policy_embedding
257-
class TestableIQLAgent(IQLAgent):
258-
def __call__(self, agent_hidden, obs, dones, pre_policy_hidden):
259-
agent_hidden, agent_embedding = jax.vmap(self.agent_rnn, in_axes=0)(
260-
agent_hidden, obs, dones
261-
)
262-
263-
pre_policy_hidden, pre_policy_embedding = jax.vmap(self.pre_policy_network, in_axes=0)(
264-
pre_policy_hidden, obs, dones
265-
)
266-
267-
# Create a mask for pre_policy_embedding
268-
mask = jnp.arange(self.num_agents) < self.num_proxy_agents
269-
mask = mask[:, None, None, None] # Reshape for broadcasting
270-
271-
# Non-proxy agents receive all-zero pre-policy embeddings
272-
pre_policy_embedding = pre_policy_embedding * mask
273-
print(pre_policy_embedding)
274-
# Apply MLP to agent embeddings before GNN
275-
gnn_input = jax.vmap(self.mlp_gnn, in_axes=0)(agent_embedding)
276-
print(f'gnn input: {gnn_input.shape}') # Expected: (num_agents, timesteps, batch_size, 14)
277-
278-
# Apply GNN to the processed agent embeddings
279-
gnn_features = jax.vmap(self.gnn, in_axes=0)(gnn_input)
280-
print(f'gnn features shape: {gnn_features.shape}') # Expected: (num_agents, timesteps, batch_size, out_embedding_dim)
281-
282-
# Concatenate embeddings for Q-value computation
283-
q_value_input = jnp.concatenate(
284-
[agent_embedding, gnn_features, pre_policy_embedding], axis=-1
285-
)
286-
q_vals = jax.vmap(self.q_value_mlp, in_axes=0)(q_value_input)
287-
288-
return agent_hidden, q_vals, pre_policy_hidden, pre_policy_embedding
289-
290-
# Initialize the testable agent
291-
testable_agent = TestableIQLAgent(
292-
action_dim=action_dim,
293-
hidden_dim=hidden_dim,
294-
num_agents=num_agents,
295-
num_proxy_agents=num_proxy_agents,
296-
init_scale=init_scale,
297-
pre_policy_output_dim=pre_policy_output_dim,
298-
pre_policy_hidden_dim=pre_policy_hidden_dim,
299-
config=config
300-
)
301-
302-
# Initialize parameters for the testable agent
303-
testable_variables = testable_agent.init(rng, agent_hidden, obs, dones, pre_policy_hidden)
304-
305-
# Apply the testable agent
306-
_, _, _, pre_policy_embedding_out = testable_agent.apply(
307-
testable_variables, agent_hidden, obs, dones, pre_policy_hidden
308-
)
309-
310-
# Create the expected mask
311-
expected_mask = jnp.arange(num_agents) < num_proxy_agents
312-
expected_mask = expected_mask[:, None, None, None] # Shape: [num_agents, 1, 1, 1]
313-
314-
# Check that pre_policy_embedding is zero for non-proxy agents
315-
non_proxy_agents = ~ (jnp.arange(num_agents) < num_proxy_agents)
316-
non_proxy_agents = non_proxy_agents[:, None, None, None] # Reshape for broadcasting
317-
318-
# Extract pre_policy_embedding for non-proxy agents
319-
masked_pre_policy = pre_policy_embedding_out * non_proxy_agents
320-
321-
# Assert that the masked pre_policy_embedding is all zeros
322-
assert jnp.all(masked_pre_policy == 0), "pre_policy_embedding is not zero for non-proxy agents"
323-
324-
print("All tests passed successfully.")
325-
326-
327-
if __name__ == "__main__":
328-
test_iql_agent()
329-
330-
331-
332-
333-
334-
335-
336169

‎baselines/QLearning/agent/pqn_agent.py‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from functools import partial
99
from agent.pre_policy_module.pre_policy_network import PrePolicyMLP
1010
from agent.gnn_module.hanabi_gnn import End2EndGCN
11-
from agent.gnn_module.hanabi_4_player_gnn import End2EndGCN4Players
11+
1212

1313

1414

@@ -77,10 +77,6 @@ def setup(self):
7777
self.gnn = nn.vmap(End2EndGCN, in_axes=0, out_axes=0,
7878
variable_axes={"params": 0},
7979
split_rngs={"params": 0})(config=self.config)
80-
elif self.num_agents == 4:
81-
self.gnn = nn.vmap(End2EndGCN4Players, in_axes=0, out_axes=0,
82-
variable_axes={"params": 0},
83-
split_rngs={"params": 0})(config=self.config)
8480
else:
8581
raise NotImplementedError("This implementation only supports 2 or 4 agents.")
8682

@@ -118,10 +114,6 @@ def __call__(self, x, train=False):
118114
q_vals: jnp.ndarray with shape (num_agents, action_dim)
119115
"""
120116

121-
# print(f"input shape: {x.shape}")
122-
# print("sum for agent 0", jnp.sum(x[0, :, -1]))
123-
#
124-
# print("sum for agent 1", jnp.sum(x[1, :, -1]))
125117
original_obs = x[:, :, :-1] # remove last
126118

127119
agent_embedding = self.agent_mlp(original_obs, train) # (num_agents, hidden_dim)
@@ -184,12 +176,6 @@ def setup(self):
184176

185177
def __call__(self, x, train=False):
186178

187-
# if self.if_augment_obs:
188-
# input = x # remove last augmented feature.
189-
# else:
190-
#
191-
# input = x[:, :, :-1]
192-
193179
agent_embedding = self.agent_mlp(x, train)
194180

195181
q_vals = self.q_network(agent_embedding)

‎baselines/QLearning/pqn_iql_pre.py‎

Lines changed: 4 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,7 @@
1818
import wandb
1919

2020
from jaxmarl import make
21-
from jaxmarl.environments.smax import map_name_to_scenario
2221
from jaxmarl.wrappers.baselines import (
23-
SMAXLogWrapper,
24-
MPELogWrapper,
2522
LogWrapper,
2623
CTRolloutManager,
2724
PrePolicyHanabiWrapper,
@@ -517,26 +514,12 @@ def _greedy_env_step(step_state, unused):
517514
return train
518515

519516

517+
520518
def env_from_config(config):
521519
env_name = config["ENV_NAME"]
522-
# smax init neeeds a scenario
523-
if "smax" in env_name.lower():
524-
config["ENV_KWARGS"]["scenario"] = map_name_to_scenario(config["MAP_NAME"])
525-
env_name = f"{config['ENV_NAME']}_{config['MAP_NAME']}"
526-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
527-
env = SMAXLogWrapper(env)
528-
# overcooked needs a layout
529-
elif "overcooked" in env_name.lower():
530-
env_name = f"{config['ENV_NAME']}_{config['ENV_KWARGS']['layout']}"
531-
config["ENV_KWARGS"]["layout"] = overcooked_layouts[
532-
config["ENV_KWARGS"]["layout"]
533-
]
534-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
535-
env = LogWrapper(env)
536-
elif "mpe" in env_name.lower():
537-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
538-
env = MPELogWrapper(env)
539-
elif "hanabi" in env_name.lower():
520+
521+
522+
if "hanabi" in env_name.lower():
540523
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
541524
env = PrePolicyHanabiWrapper(env)
542525

‎baselines/QLearning/pqn_vdn_pre.py‎

Lines changed: 4 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,17 @@
1818
import wandb
1919

2020
from jaxmarl import make
21-
from jaxmarl.environments.smax import map_name_to_scenario
2221
from jaxmarl.wrappers.baselines import (
23-
SMAXLogWrapper,
2422
MPELogWrapper,
2523
LogWrapper,
2624
CTRolloutManager,
2725
PrePolicyHanabiWrapper,
2826
)
2927

30-
from jaxmarl.wrappers.noisy_wrapper import ObsNoisyPrePolicyHanabiWrapper #NoisyPrePolicyHanabiWrapper,
3128
from jaxmarl.environments.overcooked import overcooked_layouts
3229

3330
from agent.pqn_agent import PQNAgent, BaselinePQNAgent
34-
#
35-
# jax.config.update("jax_disable_jit", True)
31+
3632

3733

3834

@@ -523,24 +519,9 @@ def _greedy_env_step(step_state, unused):
523519

524520
def env_from_config(config):
525521
env_name = config["ENV_NAME"]
526-
# smax init neeeds a scenario
527-
if "smax" in env_name.lower():
528-
config["ENV_KWARGS"]["scenario"] = map_name_to_scenario(config["MAP_NAME"])
529-
env_name = f"{config['ENV_NAME']}_{config['MAP_NAME']}"
530-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
531-
env = SMAXLogWrapper(env)
532-
# overcooked needs a layout
533-
elif "overcooked" in env_name.lower():
534-
env_name = f"{config['ENV_NAME']}_{config['ENV_KWARGS']['layout']}"
535-
config["ENV_KWARGS"]["layout"] = overcooked_layouts[
536-
config["ENV_KWARGS"]["layout"]
537-
]
538-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
539-
env = LogWrapper(env)
540-
elif "mpe" in env_name.lower():
541-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
542-
env = MPELogWrapper(env)
543-
elif "hanabi" in env_name.lower():
522+
523+
524+
if "hanabi" in env_name.lower():
544525
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
545526
env = PrePolicyHanabiWrapper(env)
546527

‎baselines/QLearning/qmix_pre.py‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,8 @@
1919
import wandb
2020

2121
from jaxmarl import make
22-
from jaxmarl.environments.smax import map_name_to_scenario
23-
from jaxmarl.environments.overcooked import overcooked_layouts
22+
2423
from jaxmarl.wrappers.baselines import (
25-
SMAXLogWrapper,
2624
MPELogWrapper,
2725
LogWrapper,
2826
CTRolloutManager,
@@ -637,31 +635,21 @@ def _greedy_env_step(step_state, unused):
637635
return train
638636

639637

638+
639+
640640
def env_from_config(config):
641641
env_name = config["ENV_NAME"]
642642
# smax init neeeds a scenario
643-
if "smax" in env_name.lower():
644-
config["ENV_KWARGS"]["scenario"] = map_name_to_scenario(config["MAP_NAME"])
645-
env_name = f"{config['ENV_NAME']}_{config['MAP_NAME']}"
646-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
647-
env = SMAXLogWrapper(env)
648-
# overcooked needs a layout
649-
elif "overcooked" in env_name.lower():
650-
env_name = f"{config['ENV_NAME']}_{config['ENV_KWARGS']['layout']}"
651-
config["ENV_KWARGS"]["layout"] = overcooked_layouts[
652-
config["ENV_KWARGS"]["layout"]
653-
]
654-
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
655-
env = LogWrapper(env)
656-
elif "mpe" in env_name.lower():
643+
if "mpe" in env_name.lower():
657644
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
658645
env = PrePolicyWrapper(env)
646+
647+
659648
else:
660649
env = make(config["ENV_NAME"], **config["ENV_KWARGS"])
661650
env = LogWrapper(env)
662651
return env, env_name
663652

664-
665653
def single_run(config):
666654

667655
config = {**config, **config["alg"]} # merge the baselines config with the main config

0 commit comments

Comments
 (0)