Refactor examples

1. Move examples dependencies out of setup.py and into examples/requirements.txt.
2. Update install-extras.sh and test-extras.sh and move to examples directory.
3. Move ./test-install.sh test code into a test (self_play_train_test.py).
4. Move tests to same directory as the units under test and name to match.
5. Add missing __init__.py files.
6. Clean up some lint errors.

PiperOrigin-RevId: 548948236
Change-Id: Ic763a4e55cb50e77041a66e79c9fb1d8ae32113a

.github/workflows/test-examples.yml

@@ -43,6 +43,6 @@ jobs:
       - name: Install Melting Pot
         run: ./install-meltingpot.sh
       - name: Install examples
-        run: ./install-extras.sh
+        run: ./examples/install.sh
       - name: Test examples
-        run: ./test-extras.sh
+        run: ./examples/test.sh

README.md

@@ -76,13 +76,12 @@ installation as follows:
 3.  Install Melting Pot:
 
     ```shell
-    pip install --editable .
+    pip install --editable .[dev]
     ```
 
 4.  (Optional) Test the installation:
 
     ```shell
-    pip install pytest-xdist
     pytest -n auto --pyargs meltingpot
     ```
 
@@ -127,7 +126,7 @@ You can try out the substrates interactively with the
 the `clean_up` substrate, you can run:
 
 ```shell
-python3 meltingpot/python/human_players/play_clean_up.py
+python meltingpot/python/human_players/play_clean_up.py
 ```
 
 You can move around with the `W`, `A`, `S`, `D` keys, Turn with `Q`, and `E`,
@@ -155,18 +154,18 @@ task suite.
 This example uses RLlib to train agents in
 self-play on a Melting Pot substrate.
 
-First you will need to install the dependencies needed by the RLlib example:
+First you will need to install the dependencies needed by the examples:
 
 ```shell
 cd <meltingpot_root>
-pip3 install -e .[rllib]
+pip install -r examples/requirements.txt
 ```
 
 Then you can run the training experiment using:
 
 ```shell
-cd <meltingpot_root>/examples/rllib
-python3 self_play_train.py
+cd examples/rllib
+python self_play_train.py
 ```
 
 #### PettingZoo and Stable-Baselines3
@@ -179,12 +178,9 @@ can be found [here](examples/pettingzoo/utils.py).
 
 ```shell
 cd <meltingpot_root>
-pip3 install -e .[pettingzoo]
-```
-
-```shell
-cd <meltingpot_root>/examples/pettingzoo
-python3 sb3_train.py
+pip install -r examples/requirements.txt
+cd examples/pettingzoo
+python sb3_train.py
 ```
 
 ## Documentation

examples/gym/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

install-extras.sh → examples/install.sh

@@ -13,27 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Installs meltingpot extras on Linux/macOS.
+# Installs meltingpot examples.
 
 set -euxo pipefail
 
+cd "$(dirname "$0")"
 
-function check_setup() {
-  echo -e "\nChecking meltingpot is installed..."
-  python -c 'import meltingpot'
-}
+pip install --upgrade pip
+pip install setuptools
 
-
-function install_extras() {
-  echo -e "\nInstalling meltingpot extras..."
-  pip install .[rllib,pettingzoo]
-}
-
-
-function main() {
-  check_setup
-  install_extras
-}
-
-
-main "$@"
+pip install -r requirements.txt

examples/pettingzoo/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

examples/pettingzoo/utils.py

@@ -22,7 +22,7 @@
 from pettingzoo import utils as pettingzoo_utils
 from pettingzoo.utils import wrappers
 
-from .. import utils
+from ..gym import utils
 
 PLAYER_STR_FORMAT = 'player_{index}'
 MAX_CYCLES = 1000

examples/requirements.txt

@@ -0,0 +1,13 @@
+absl_py
+dm_env
+dm_meltingpot
+dm_tree
+dmlab2d
+gym
+matplotlib
+ml_collections
+numpy<1.23  # Needed by Ray because it uses `np.bool`.
+ray[rllib,default]==2.0.0
+stable_baselines3
+supersuit>=3.7.2
+torch

examples/rllib/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

examples/rllib/self_play_train.py

@@ -133,23 +133,31 @@ def policy_mapping_fn(agent_id, **kwargs):
   return config
 
 
-def main():
+def train(config, num_iterations=1):
+  """Trains a model.
 
-  config = get_config()
-  tune.register_env("meltingpot", utils.env_creator)
+  Args:
+    config: model config
+    num_iterations: number of iterations ot train for.
 
-  # 6. Initialize ray, train and save
+  Returns:
+    Training results.
+  """
+  tune.register_env("meltingpot", utils.env_creator)
   ray.init()
-
   stop = {
-      "training_iteration": 1,
+      "training_iteration": num_iterations,
   }
-
-  results = tune.Tuner(
+  return tune.Tuner(
       "PPO",
       param_space=config.to_dict(),
       run_config=air.RunConfig(stop=stop, verbose=1),
   ).fit()
+
+
+def main():
+  config = get_config()
+  results = train(config, num_iterations=1)
   print(results)
   assert results.num_errors == 0
 

examples/rllib/self_play_train_test.py

@@ -0,0 +1,34 @@
+# Copyright 2020 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the self_play_train.py."""
+
+from absl.testing import absltest
+
+from . import self_play_train
+
+
+class TrainingTests(absltest.TestCase):
+  """Tests for MeltingPotEnv for RLLib."""
+
+  def test_training(self):
+    config = self_play_train.get_config(
+        num_rollout_workers=1,
+        rollout_fragment_length=10,
+        train_batch_size=20,
+        sgd_minibatch_size=20,
+        fcnet_hiddens=(4,),
+        post_fcnet_hiddens=(4,),
+        lstm_cell_size=2)
+    results = self_play_train.train(config, num_iterations=1)
+    self.assertEqual(results.num_errors, 0)

examples/rllib/utils.py

@@ -24,9 +24,9 @@
 import numpy as np
 from ray.rllib import algorithms
 from ray.rllib.env import multi_agent_env
-from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.policy import sample_batch
 
-from .. import utils
+from ..gym import utils
 
 
 PLAYER_STR_FORMAT = 'player_{index}'
@@ -64,9 +64,9 @@ def reset(self):
     timestep = self._env.reset()
     return utils.timestep_to_observations(timestep)
 
-  def step(self, action):
+  def step(self, action_dict):
     """See base class."""
-    actions = [action[agent_id] for agent_id in self._ordered_agent_ids]
+    actions = [action_dict[agent_id] for agent_id in self._ordered_agent_ids]
     timestep = self._env.step(actions)
     rewards = {
         agent_id: timestep.reward[index]
@@ -147,7 +147,7 @@ class RayModelPolicy(policy.Policy):
 
   def __init__(self,
                model: algorithms.Algorithm,
-               policy_id: str = DEFAULT_POLICY_ID) -> None:
+               policy_id: str = sample_batch.DEFAULT_POLICY_ID) -> None:
     """Initialize a policy instance.
 
     Args:

examples/rllib/tests/test_environment.py → examples/rllib/utils_test.py

@@ -11,14 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""MeltingPotEnv as a MultiAgentEnv wrapper to interface with RLLib."""
+"""Tests for utils.py."""
 
 from absl.testing import absltest
 from gym.spaces import discrete
 from meltingpot import substrate
 from meltingpot.configs.substrates import commons_harvest_open
 
-from .. import utils
+from . import utils
 
 
 class MeltingPotEnvTests(absltest.TestCase):

examples/test.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Copyright 2020 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Tests meltingpot examples.
+
+set -euxo pipefail
+
+cd "$(dirname "$0")"
+
+pip install --upgrade pip
+pip install pytest-xdist
+
+pytest

setup.py

@@ -135,20 +135,5 @@ def build_assets(self):
             'pytest-xdist',
             'pytype',
         ],
-        # Required for rllib example.
-        'rllib': [
-            'gym',
-            'ray[rllib,default]==2.0.0',
-            'numpy<1.23',  # Needed by Ray because it uses `np.bool`.
-        ],
-        # Required for pettingzoo example.
-        'pettingzoo': [
-            'gym',
-            'matplotlib',
-            'pettingzoo>=1.22.3',
-            'stable-baselines3',
-            'supersuit>=3.7.2',
-            'torch',
-        ],
     },
 )