exp: pettingzoo.

Dockerfile

@@ -42,9 +42,9 @@ RUN pip install --quiet --upgrade pip setuptools wheel &&  \
 # ENV SUPPRESS_GR_PROMPT 1
 # RUN ./install_environments/mamujoco.sh
 
-# RUN ./install_environments/pettingzoo.sh
+RUN ./install_environments/pettingzoo.sh
 
-RUN ./install_environments/flatland.sh
+# RUN ./install_environments/flatland.sh
 
 # Copy all code
 COPY ./examples ./examples

install_environments/requirements/pettingzoo.txt

@@ -2,7 +2,7 @@ autorom
 gym
 numpy
 opencv-python
-pettingzoo==1.22.3
+pettingzoo==1.23.1
 pygame
 pymunk
 scipy

manifest.yaml

@@ -13,19 +13,19 @@ builder:
 spec:
   operator: tf
   image: og-marl
-  # command: python baselines/main.py --env=pettingzoo --scenario=pursuit
-  command: wandb agent off-the-grid-marl-team/og-marl-baselines/i87akim6
+  command: python baselines/main.py --env=pettingzoo --scenario=pursuit
+  # command: wandb agent off-the-grid-marl-team/og-marl-baselines/i87akim6
 
   types:
     Worker:
       count: 60
 
       resources:
         # 10-15 cpus for single process
-        cpus: 4
+        cpus: 8
         # mandatory (must be >= 2): RAM ratio wanted (in GB)
         # memory = cpus * ramRatio
-        ramRatio: 2
+        ramRatio: 3
         accelerators:
           gpu:
             count: 0

og_marl/replay_buffers.py

@@ -97,14 +97,13 @@ def populate_from_vault(
             ).read()
 
             # Recreate the buffer and associated pure functions
-            self._max_size = self._buffer_state.current_index
             self._replay_buffer = fbx.make_trajectory_buffer(
                 add_batch_size=1,
                 sample_batch_size=self._batch_size,
                 sample_sequence_length=self._sequence_length,
                 period=1,
                 min_length_time_axis=1,
-                max_size=self._max_size,
+                max_size=self._sequence_length,
             )
             self._buffer_sample_fn = jax.jit(self._replay_buffer.sample)
             self._buffer_add_fn = jax.jit(self._replay_buffer.add)