explicit the jaxlib version to prevent jax-jaxlib mismatches + cleanup

FLAIROx · Mar 22, 2024 · a8dab70 · a8dab70
1 parent 491bbf7
commit a8dab70
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 38 deletions.
diff --git a/jaxmarl/environments/hanabi/README.md b/jaxmarl/environments/hanabi/README.md
@@ -142,11 +142,11 @@ Other Hand:
 3 Card: G1, Hints: , Possible: RYGWB12345, Belief: [R1: 0.060 Y1: 0.060 G1: 0.060 W1: 0.060 B1: 0.060]
 4 Card: Y2, Hints: , Possible: RYGWB12345, Belief: [R1: 0.060 Y1: 0.060 G1: 0.060 W1: 0.060 B1: 0.060]
 Your Hand:
-0 Hints: , Possible: RYGWB2345, Belief: [R2: 0.057 R3: 0.057 R4: 0.057 Y2: 0.057 Y3: 0.057]
-1 Hints: 1, Possible: RYGWB1, Belief: [R1: 0.200 Y1: 0.200 G1: 0.200 W1: 0.200 B1: 0.200]
-2 Hints: 1, Possible: RYGWB1, Belief: [R1: 0.200 Y1: 0.200 G1: 0.200 W1: 0.200 B1: 0.200]
-3 Hints: , Possible: RYGWB2345, Belief: [R2: 0.057 R3: 0.057 R4: 0.057 Y2: 0.057 Y3: 0.057]
-4 Hints: , Possible: RYGWB2345, Belief: [R2: 0.057 R3: 0.057 R4: 0.057 Y2: 0.057 Y3: 0.057]
+0 Hints: ,  Possible: RYGWB2345, Belief: [R2: 0.057 R3: 0.057 R4: 0.057 Y2: 0.057 Y3: 0.057]
+1 Hints: 1, Possible: RYGWB1,    Belief: [R1: 0.200 Y1: 0.200 G1: 0.200 W1: 0.200 B1: 0.200]
+2 Hints: 1, Possible: RYGWB1,    Belief: [R1: 0.200 Y1: 0.200 G1: 0.200 W1: 0.200 B1: 0.200]
+3 Hints: ,  Possible: RYGWB2345, Belief: [R2: 0.057 R3: 0.057 R4: 0.057 Y2: 0.057 Y3: 0.057]
+4 Hints: ,  Possible: RYGWB2345, Belief: [R2: 0.057 R3: 0.057 R4: 0.057 Y2: 0.057 Y3: 0.057]
 Last action: H1
 Cards afected: [1 2]
 Legal Actions: ['D0', 'D1', 'D2', 'D3', 'D4', 'P0', 'P1', 'P2', 'P3', 'P4', 'HY', 'HG', 'HW', 'H1', 'H2', 'H3', 'H4', 'H5']
@@ -161,18 +161,16 @@ python manual_game.py \
   --player0 "manual" \
   --player1 "obl" \
   --weight1 "./pretrained/obl-r2d2-flax/icml_OBL1/OFF_BELIEF1_SHUFFLE_COLOR0_BZA0_BELIEF_a.safetensors" \
-  --seed 7
 ```
 
 Or to look an obl model playing with itself:
 
 ```
 python manual_game.py \
-  --player0 "manual" \
+  --player0 "obl" \
   --player1 "obl" \
   --weight0 "./pretrained/obl-r2d2-flax/icml_OBL1/OFF_BELIEF1_SHUFFLE_COLOR0_BZA0_BELIEF_a.safetensors" \
   --weight1 "./pretrained/obl-r2d2-flax/icml_OBL1/OFF_BELIEF1_SHUFFLE_COLOR0_BZA0_BELIEF_a.safetensors" \
-  --seed 7
 ```
 
 ## Citation

diff --git a/jaxmarl/environments/hanabi/pretrained/decks_test.json b/jaxmarl/environments/hanabi/pretrained/decks_test.json
diff --git a/jaxmarl/environments/hanabi/render_obs.py b/jaxmarl/environments/hanabi/render_obs.py
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -1,5 +1,6 @@
 # requirements are alligned with nvcr.io/nvidia/jax:23.10-py3 image
 jax==0.4.17
+jaxlib==0.4.17
 flax==0.7.4
 chex==0.1.84
 optax==0.1.7

diff --git a/tests/hanabi/test_hanabi.py b/tests/hanabi/test_hanabi.py
@@ -89,6 +89,7 @@ def test_injected_decks():
     This tests consists in injecting in the Hanabi environment a set of decks and actions that are known to produce a certain score.
     The test checks if the scores produced by the environment are the same as the expected ones.
     """
+    print('Hanabi Test: test_injected_decks')
     actions_seq = get_action_sequences()
     decks = get_decks()
     true_scores = get_scores()

diff --git a/tests/switch_riddle/test_heuristic.py b/tests/switch_riddle/test_heuristic.py
@@ -13,7 +13,7 @@ def action_for_agent_in_room(env, state, action_name):
     return actions
 
 @pytest.mark.parametrize("seed", [0, 999])
-@pytest.mark.parametrize("num_agents", [3, 5, 10, 50, 100])
+@pytest.mark.parametrize("num_agents", [3, 5, 10])
 class TestHeuristic:
 
     def test_light_switch(self, seed, num_agents):