Merge pull request #1175 from kubicon:master

PiperOrigin-RevId: 625293351
Change-Id: Idc4c2c330119a89bffb176ab5a4e03d0b9749807

open_spiel/python/CMakeLists.txt

@@ -267,6 +267,7 @@ if (OPEN_SPIEL_ENABLE_JAX)
     jax/nfsp_jax_test.py
     jax/opponent_shaping_jax_test.py
     jax/policy_gradient_jax_test.py
+    jax/cfr/jax_cfr_test.py
     algorithms/rnad/rnad_test.py
     coalitional_games/least_core_lagrangian_test.py
     mfg/algorithms/fictitious_play_test.py

open_spiel/python/jax/cfr/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2019 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
+#
+#      http://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.

open_spiel/python/jax/cfr/compare_cfr_with_jax.py

@@ -0,0 +1,109 @@
+# Copyright 2019 DeepMind Technologies Ltd. All rights reserved.
+#
+# 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
+#
+#     http://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.
+
+"""This compares the speed/results of Jax CFR to of the original impl of CFR.
+
+The results slightly differ due to different rounding of regrets between
+original implmentation and CFR. When setting clamping of regrets to 1e-8 the
+results are exactly the same.
+"""
+
+
+# pylint: disable=g-importing-member
+
+import time
+from open_spiel.python.algorithms.best_response import BestResponsePolicy
+from open_spiel.python.algorithms.cfr import CFRPlusSolver
+from open_spiel.python.jax.cfr.jax_cfr import JaxCFR
+import pyspiel
+
+
+def compare_cfr_with_jax_cfr(game):
+  """Do the comparison."""
+
+  start = time.time()
+  jax_cfr = JaxCFR(game)
+  print(time.time() - start)
+  jax_cfr.multiple_steps(10000)
+  print(time.time() - start)
+
+  start = time.time()
+  print(time.time() - start)
+  cfr = CFRPlusSolver(game)
+  for _ in range(1000):
+    cfr.evaluate_and_update_policy()
+
+  print(time.time() - start)
+
+  jax_strat = jax_cfr.average_policy()
+  jax_br1 = BestResponsePolicy(jax_cfr.game, 1, jax_strat)
+  jax_br2 = BestResponsePolicy(jax_cfr.game, 0, jax_strat)
+
+  cfr_strat = jax_cfr.average_policy()
+  cfr_br1 = BestResponsePolicy(jax_cfr.game, 1, cfr_strat)
+  cfr_br2 = BestResponsePolicy(jax_cfr.game, 0, cfr_strat)
+
+  print("Jax P1: ", jax_br1.value(jax_cfr.game.new_initial_state()))
+  print("CFR P1: ", cfr_br1.value(jax_cfr.game.new_initial_state()))
+  print("Jax P2: ", jax_br2.value(jax_cfr.game.new_initial_state()))
+  print("CFR P2: ", cfr_br2.value(jax_cfr.game.new_initial_state()))
+
+
+# Speed Results:
+#   Original: 139.60753107070923
+#   Jax CPU: 3.7404067516326904
+def compare_leduc():
+  game = pyspiel.load_game("leduc_poker")
+  compare_cfr_with_jax_cfr(game)
+
+
+# Speed Results:
+#   Original: 335.6707363128662
+#   Jax CPU: 7.59996485710144
+def compare_battleship():
+  game_params = {
+      "board_height": 2,
+      "board_width": 2,
+      "num_shots": 4,
+      "ship_sizes": "[2]",
+      "ship_values": "[1]",
+      "allow_repeated_shots": False,
+  }
+  game = pyspiel.load_game("battleship", game_params)
+  compare_cfr_with_jax_cfr(game)
+
+
+# Speed Results:
+#   Original: 14.667663097381592
+#   Jax CPU: 1.068636417388916
+def compare_goofspiel_descending():
+  game_params = {"num_cards": 4, "imp_info": True, "points_order": "descending"}
+  game = pyspiel.load_game_as_turn_based("goofspiel", game_params)
+  compare_cfr_with_jax_cfr(game)
+
+
+# Speed Results:
+#   Original: 6.639796733856201
+#   Jax CPU: 0.8599820137023926
+def compare_goofspiel_randomized():
+  game_params = {"num_cards": 3, "imp_info": True, "points_order": "random"}
+  game = pyspiel.load_game_as_turn_based("goofspiel", game_params)
+  compare_cfr_with_jax_cfr(game)
+
+
+if __name__ == "__main__":
+  compare_leduc()
+  compare_battleship()
+  compare_goofspiel_descending()
+  compare_goofspiel_randomized()