update

connectfour/app.py

@@ -2,25 +2,25 @@ import time
 
 import gradio as gr
 import numpy as np
-import ray
-import ray.rllib.algorithms.ppo as ppo
 from pettingzoo.classic import connect_four_v3
+from ray.rllib.utils.framework import try_import_torch
 from ray.tune import register_env
 
-from connectfour.checkpoint import CHECKPOINT
-from connectfour.training.models import Connect4MaskModel
 from connectfour.training.wrappers import Connect4Env
+from models import MODEL_PATH
+import onnxruntime as ort
+from ray.rllib.algorithms.algorithm import Algorithm
+from connectfour.checkpoint import CHECKPOINT
 
-POLICY_ID = "learned_v5"
+torch, nn = try_import_torch()
 
 # poetry export -f requirements.txt --output requirements.txt --without-hashes
+# tensorboard --logdir ~/ray_results/
 # gradio connectfour/app.py
 
 
 class Connect4:
     def __init__(self, who_plays_first) -> None:
-        ray.init(include_dashboard=False, ignore_reinit_error=True)
-
         # define how to make the environment
         env_creator = lambda config: connect_four_v3.env(render_mode="rgb_array")
 
@@ -44,25 +44,46 @@ class Connect4:
 
         return self.render_and_state
 
-    def get_algo(self, checkpoint):
-        config = (
-            ppo.PPOConfig()
-            .environment("connect4")
-            .framework("torch")
-            .training(model={"custom_model": Connect4MaskModel})
+    def get_algo(self):
+        # self.pytorch_model = torch.load(MODEL_PATH / "model.pt")
+        # self.algo = Algorithm.from_checkpoint(checkpoint=CHECKPOINT)
+        self.session = ort.InferenceSession(str(MODEL_PATH / "model.onnx"), None)
+
+    def compute_action(self, obs):
+        return self.pytorch_model(
+            input_dict={"obs": self.flatten_obs(obs)},
+        )
+
+    def flatten_obs(self, obs):
+        flatten_action_mask = torch.from_numpy(obs["action_mask"])
+        flatten_observation = torch.flatten(
+            torch.from_numpy(obs["observation"]), end_dim=2
         )
-        config.explore = False
-        self.algo = config.build()
-        self.algo.restore(checkpoint)
+        flatten_obs = torch.concat([flatten_action_mask, flatten_observation])
+        return flatten_obs[None, :]
 
     def play(self, action=None):
         if self.has_erroneous_state():
             return self.blue_screen()
 
         if self.human != self.player_id:
-            action = self.algo.compute_single_action(
-                self.obs[self.player_id], policy_id=POLICY_ID
+            # action = self.algo.compute_single_action(
+            #     self.obs[self.player_id], policy_id="learned_v9"
+            # )
+            # Torch
+            # action = self.compute_action(self.obs[self.player_id])
+            # action = int(torch.argmax(action[0]))
+            # ONNX
+            action = self.session.run(
+                ["output"],
+                {
+                    "obs": self.flatten_obs(self.obs[self.player_id])
+                    .numpy()
+                    .astype(np.float32),
+                    "state_ins": [],
+                },
             )
+            action = int(np.argmax(action[0]))
 
         if action not in self.legal_moves:
             action = np.random.choice(self.legal_moves)
@@ -114,7 +135,7 @@ demo = gr.Blocks()
 
 with demo:
     connect4 = Connect4("You")
-    connect4.get_algo(str(CHECKPOINT))
+    connect4.get_algo()
 
     with gr.Row():
         with gr.Column(scale=1):

connectfour/checkpoint/policies/always_same/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d278413093ad1bc4f227279e3dab7be04ebd70ca1ed156a1363515c69d0a858e
-size 10992

connectfour/checkpoint/policies/always_same/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/beat_last/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cd422258c16de0866599730a5a5b2b48e2ee81cbae69f9d5471deeae76c42b47
-size 10992

connectfour/checkpoint/policies/beat_last/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/learned/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2a517583e5fcad7e483bca619723583cc6928499390c1fcfc25d907e109cd4b4
-size 2139442

connectfour/checkpoint/policies/learned/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/learned_v1/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:276c26007c2419a688c27f9dfa70c20fecb468a0aa07d28d6a9e8099bbc849be
-size 2139439

connectfour/checkpoint/policies/learned_v1/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/learned_v2/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e37a485d3a54f7a8b194693e7a61f790e67071358130178fa01cdbd840c4a4da
-size 2139439

connectfour/checkpoint/policies/learned_v2/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/learned_v3/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9f90899ae98a387e312333b234041c68b9c50da4af92ee5250686087a39eebb3
-size 2139439

connectfour/checkpoint/policies/learned_v3/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/learned_v4/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3af3b3fe41bac489cb693af387b1ccc4437a532a78d539b3abb4cc5f77929592
-size 2139439

connectfour/checkpoint/policies/learned_v4/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/learned_v5/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f2b28b979e2f4411d196e03ca75ea7f25f7601bb997aa8bcdcf1d49c9ea30754
-size 2139439

connectfour/checkpoint/policies/learned_v5/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/linear/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4f70d44ac661632dc0557204abe34308dfb25b800a668b49c2efd9a2a73a7bc0
-size 10992

connectfour/checkpoint/policies/linear/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/policies/random/policy_state.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f3b1ab86bada035779feedb2b92ae0a64f6d9474bb4f0ae44324e17d65659764
-size 10992

connectfour/checkpoint/policies/random/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Policy", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/checkpoint/rllib_checkpoint.json

@@ -1 +0,0 @@
-{"type": "Algorithm", "checkpoint_version": "1.0", "ray_version": "2.3.1", "ray_commit": "5f14cee8dfc6d61ec4fd3bc2c440f9944e92b33a"}

connectfour/training/callbacks.py

@@ -1,33 +1,44 @@
-from ray.rllib.algorithms.callbacks import DefaultCallbacks
+from collections import deque
+
 import numpy as np
+from ray.rllib.algorithms.callbacks import DefaultCallbacks
+
+from connectfour.training.dummy_policies import (
+    AlwaysSameHeuristic,
+    BeatLastHeuristic,
+    LinearHeuristic,
+    RandomHeuristic,
+)
 
 
-def create_self_play_callback(win_rate_thr, opponent_policies):
+def create_self_play_callback(win_rate_thr, opponent_policies, opponent_count=10):
     class SelfPlayCallback(DefaultCallbacks):
         win_rate_threshold = win_rate_thr
 
         def __init__(self):
             super().__init__()
             self.current_opponent = 0
+            self.opponent_policies = deque(opponent_policies, maxlen=opponent_count)
+            self.policy_to_remove = None
+            self.frozen_policies = {
+                "always_same": AlwaysSameHeuristic,
+                "linear": LinearHeuristic,
+                "beat_last": BeatLastHeuristic,
+                "random": RandomHeuristic,
+            }
 
         def on_train_result(self, *, algorithm, result, **kwargs):
-            # Get the win rate for the train batch.
-            # Note that normally, one should set up a proper evaluation config,
-            # such that evaluation always happens on the already updated policy,
-            # instead of on the already used train_batch.
+            """Called at the end of Algorithm.train().
+
+            Args:
+                algorithm: Current Algorithm instance.
+                result: Dict of results returned from Algorithm.train() call.
+                    You can mutate this object to add additional metrics.
+                kwargs: Forward compatibility placeholder.
+            """
             main_rew = result["hist_stats"].pop("policy_learned_reward")
             opponent_rew = result["hist_stats"].pop("episode_reward")
 
-            if len(main_rew) != len(opponent_rew):
-                raise Exception(
-                    "len(main_rew) != len(opponent_rew)",
-                    len(main_rew),
-                    len(opponent_rew),
-                    result["hist_stats"].keys(),
-                    "episode len",
-                    len(opponent_rew),
-                )
-
             won = 0
             for r_main, r_opponent in zip(main_rew, opponent_rew):
                 if r_main > r_opponent:
@@ -35,54 +46,85 @@ def create_self_play_callback(win_rate_thr, opponent_policies):
             win_rate = won / len(main_rew)
 
             result["win_rate"] = win_rate
-            print(f"Iter={algorithm.iteration} win-rate={win_rate} -> ", end="")
+            print(f"Iter={algorithm.iteration} win-rate={win_rate}")
 
-            # If win rate is good -> Snapshot current policy and play against
-            # it next, keeping the snapshot fixed and only improving the "learned"
-            # policy.
             if win_rate > self.win_rate_threshold:
-                self.current_opponent += 1
-                new_pol_id = f"learned_v{self.current_opponent}"
-                print(
-                    f"Iter={algorithm.iteration} ### Adding new opponent to the mix ({new_pol_id})."
-                )
+                if len(self.opponent_policies) == self.opponent_policies.maxlen:
+                    self.policy_to_remove = self.opponent_policies[0]
+
+                new_pol_id = None
+                while new_pol_id is None:
+                    if np.random.choice(range(6)) == 0:
+                        new_pol_id = np.random.choice(list(self.frozen_policies.keys()))
+                    else:
+                        self.current_opponent += 1
+                        new_pol_id = f"learned_v{self.current_opponent}"
+
+                    if new_pol_id in self.opponent_policies:
+                        new_pol_id = None
+                    else:
+                        self.opponent_policies.append(new_pol_id)
+
+                print("Non trainable policies", list(self.opponent_policies))
 
-                # Re-define the mapping function, such that "learned" is forced
-                # to play against any of the previously played policies
-                # (excluding "random").
                 def policy_mapping_fn(agent_id, episode, worker, **kwargs):
-                    # agent_id = [0|1] -> policy depends on episode ID
-                    # This way, we make sure that both policies sometimes play
-                    # (start player) and sometimes agent1 (player to move 2nd).
                     return (
                         "learned"
                         if episode.episode_id % 2 == int(agent_id[-1:])
-                        else np.random.choice(
-                            opponent_policies
-                            + [
-                                f"learned_v{i}"
-                                for i in range(1, self.current_opponent + 1)
-                            ]
-                        )
+                        else np.random.choice(list(self.opponent_policies))
                     )
 
-                new_policy = algorithm.add_policy(
-                    policy_id=new_pol_id,
-                    policy_cls=type(algorithm.get_policy("learned")),
-                    policy_mapping_fn=policy_mapping_fn,
+                print(
+                    f"Iter={algorithm.iteration} Adding new opponent to the mix ({new_pol_id}). League size {len(self.opponent_policies) + 1}"
                 )
 
-                # Set the weights of the new policy to the learned policy.
-                # We'll keep training the learned policy, whereas `new_pol_id` will
-                # remain fixed.
-                learned_state = algorithm.get_policy("learned").get_state()
-                new_policy.set_state(learned_state)
-                # We need to sync the just copied local weights (from learned policy)
-                # to all the remote workers as well.
+                if new_pol_id in list(self.frozen_policies.keys()):
+                    new_policy = algorithm.add_policy(
+                        policy_id=new_pol_id,
+                        policy_cls=self.frozen_policies[new_pol_id],
+                        policy_mapping_fn=policy_mapping_fn,
+                    )
+                else:
+                    new_policy = algorithm.add_policy(
+                        policy_id=new_pol_id,
+                        policy_cls=type(algorithm.get_policy("learned")),
+                        policy_mapping_fn=policy_mapping_fn,
+                    )
+                    learned_state = algorithm.get_policy("learned").get_state()
+                    new_policy.set_state(learned_state)
                 algorithm.workers.sync_weights()
+
             else:
-                print("not good enough; will keep learning ...")
+                print("Not good enough... Keep learning ...")
+
+            result["league_size"] = len(self.opponent_policies) + 1
+
+        def on_evaluate_end(self, *, algorithm, evaluation_metrics, **kwargs):
+            """Runs when the evaluation is done.
+
+            Runs at the end of Algorithm.evaluate().
 
-            result["league_size"] = self.current_opponent + len(opponent_policies) + 1
+            Args:
+                algorithm: Reference to the algorithm instance.
+                evaluation_metrics: Results dict to be returned from algorithm.evaluate().
+                    You can mutate this object to add additional metrics.
+                kwargs: Forward compatibility placeholder.
+            """
+
+            def policy_mapping_fn(agent_id, episode, worker, **kwargs):
+                return (
+                    "learned"
+                    if episode.episode_id % 2 == int(agent_id[-1:])
+                    else np.random.choice(list(self.opponent_policies))
+                )
+
+            if self.policy_to_remove is not None:
+                print("Remove ", self.policy_to_remove, "from opponent policies")
+                algorithm.remove_policy(
+                    self.policy_to_remove,
+                    policy_mapping_fn=policy_mapping_fn,
+                )
+                self.policy_to_remove = None
+                algorithm.workers.sync_weights()
 
     return SelfPlayCallback

connectfour/training/dummy_policies.py

@@ -23,6 +23,10 @@ class HeuristicBase(Policy):
         """No weights to set."""
         pass
 
+    @override(Policy)
+    def export_model(self, export_dir: str, onnx=None) -> None:
+        pass
+
     @override(Policy)
     def compute_actions(
         self,

connectfour/training/train.py

@@ -8,6 +8,7 @@ from ray import air, tune
 from ray.rllib.policy.policy import PolicySpec
 from ray.rllib.utils.framework import try_import_torch
 from ray.tune import CLIReporter, register_env
+from ray.rllib.algorithms.algorithm import Algorithm
 
 from connectfour.training.callbacks import create_self_play_callback
 from connectfour.training.dummy_policies import (
@@ -29,12 +30,21 @@ def get_cli_args():
     python connectfour/training/train.py --num-cpus 4 --num-gpus 1 --stop-iters 10 --win-rate-threshold 0.50
     python connectfour/training/train.py --num-gpus 1 --stop-iters 1 --win-rate-threshold 0.50
     python connectfour/training/train.py --num-cpus 5 --num-gpus 1 --stop-iters 200
+    python connectfour/training/train.py --num-cpus 5 --num-gpus 1 --win-rate-threshold 0.95 --stop-iters 2000 > training.log 2>&1
+    python connectfour/training/train.py --num-cpus 5 --num-gpus 1 --win-rate-threshold 0.96 --stop-iters 10000 > training.log 2>&1
+    python connectfour/training/train.py --num-cpus 5 --num-gpus 1 --win-rate-threshold 0.99 --stop-iters 5000 --from-checkpoint ~/ray_results/PPO/PPO_connect4_8414a_00000_0_2023-04-03_12-44-31/checkpoint_004000 > training.log 2>&1
     """
     parser = argparse.ArgumentParser()
     parser.add_argument("--num-cpus", type=int, default=0)
     parser.add_argument("--num-gpus", type=int, default=0)
     parser.add_argument("--num-workers", type=int, default=2)
-
+    parser.add_argument(
+        "--from-checkpoint",
+        type=str,
+        default=None,
+        help="Full path to a experiment directory to resume tuning from "
+        "a previously saved Algorithm state.",
+    )
     parser.add_argument(
         "--stop-iters", type=int, default=200, help="Number of iterations to train."
     )
@@ -57,13 +67,6 @@ def get_cli_args():
     return args
 
 
-def select_policy(agent_id, episode, **kwargs):
-    if episode.episode_id % 2 == int(agent_id[-1:]):
-        return "learned"
-    else:
-        return random.choice(["always_same", "beat_last", "random", "linear"])
-
-
 if __name__ == "__main__":
     args = get_cli_args()
 
@@ -80,20 +83,23 @@ if __name__ == "__main__":
     # register that way to make the environment under an rllib name
     register_env("connect4", lambda config: Connect4Env(env_creator(config)))
 
+    def select_policy(agent_id, episode, **kwargs):
+        if episode.episode_id % 2 == int(agent_id[-1:]):
+            return "learned"
+        else:
+            return random.choice(["always_same", "beat_last", "random", "linear"])
+
     config = (
-        ppo.PPOConfig()
-        .environment("connect4")
-        .framework("torch")
-        .training(model={"custom_model": Connect4MaskModel})
-        .callbacks(
-            create_self_play_callback(
-                win_rate_thr=args.win_rate_threshold,
-                opponent_policies=["always_same", "beat_last", "random", "linear"],
+        (
+            ppo.PPOConfig()
+            .environment("connect4")
+            .framework("torch")
+            .training(model={"custom_model": Connect4MaskModel})
+            .rollouts(
+                num_rollout_workers=args.num_workers,
+                num_envs_per_worker=5,
             )
-        )
-        .rollouts(
-            num_rollout_workers=args.num_workers,
-            num_envs_per_worker=5,
+            # .checkpointing(checkpoint_trainable_policies_only=True)
         )
         .multi_agent(
             policies={
@@ -106,19 +112,88 @@ if __name__ == "__main__":
             policy_mapping_fn=select_policy,
             policies_to_train=["learned"],
         )
+        .callbacks(
+            create_self_play_callback(
+                win_rate_thr=args.win_rate_threshold,
+                opponent_policies=["always_same", "beat_last", "random", "linear"],
+                opponent_count=15,
+            )
+        )
+        .evaluation(evaluation_interval=1)
     )
 
-    stop = {
-        "timesteps_total": args.stop_timesteps,
-        "training_iteration": args.stop_iters,
-    }
-
-    results = tune.Tuner(
-        "PPO",
-        param_space=config.to_dict(),
-        run_config=air.RunConfig(
-            stop=stop,
-            verbose=2,
+    if args.from_checkpoint is None:
+        stop = {
+            "timesteps_total": args.stop_timesteps,
+            "training_iteration": args.stop_iters,
+        }
+
+        results = tune.Tuner(
+            "PPO",
+            param_space=config.to_dict(),
+            run_config=air.RunConfig(
+                stop=stop,
+                verbose=2,
+                progress_reporter=CLIReporter(
+                    metric_columns={
+                        "training_iteration": "iter",
+                        "time_total_s": "time_total_s",
+                        "timesteps_total": "ts",
+                        "episodes_this_iter": "train_episodes",
+                        "policy_reward_mean/learned": "reward",
+                        "win_rate": "win_rate",
+                        "league_size": "league_size",
+                    },
+                    mode="max",
+                    metric="win_rate",
+                    sort_by_metric=True,
+                ),
+                checkpoint_config=air.CheckpointConfig(
+                    num_to_keep=10,
+                    checkpoint_at_end=True,
+                    checkpoint_frequency=10,
+                    checkpoint_score_order="max",
+                ),
+            ),
+        ).fit()
+
+        best_checkpoint = results.get_best_result(
+            metric="win_rate", mode="max"
+        ).checkpoint
+        print("Best checkpoint", best_checkpoint)
+
+    else:
+        algo = Algorithm.from_checkpoint(checkpoint=args.from_checkpoint)
+
+        config = algo.config.copy(False)
+        config.checkpointing(export_native_model_files=True)
+
+        opponent_policies = list(algo.workers.local_worker().policy_map.keys())
+        opponent_policies.remove("learned")
+        opponent_policies.sort()
+
+        config.callbacks(
+            create_self_play_callback(
+                win_rate_thr=args.win_rate_threshold,
+                opponent_policies=opponent_policies,
+                opponent_count=len(opponent_policies),
+            )
+        )
+        config.evaluation(evaluation_interval=None)
+
+        analysis = tune.run(
+            "PPO",
+            config=config.to_dict(),
+            restore=args.from_checkpoint,
+            checkpoint_freq=10,
+            checkpoint_at_end=True,
+            keep_checkpoints_num=10,
+            mode="max",
+            metric="win_rate",
+            stop={
+                "win_rate": args.win_rate_threshold,
+                "training_iteration": args.stop_iters,
+            },
             progress_reporter=CLIReporter(
                 metric_columns={
                     "training_iteration": "iter",
@@ -128,21 +203,18 @@ if __name__ == "__main__":
                     "policy_reward_mean/learned": "reward",
                     "win_rate": "win_rate",
                     "league_size": "league_size",
-                },
-                mode="max",
-                metric="win_rate",
-                sort_by_metric=True,
-            ),
-            checkpoint_config=air.CheckpointConfig(
-                checkpoint_at_end=True,
-                checkpoint_frequency=10,
+                }
             ),
-        ),
-    ).fit()
+        )
 
-    print(
-        "Best checkpoint",
-        results.get_best_result(metric="win_rate", mode="max").checkpoint,
-    )
+        algo = Algorithm.from_checkpoint(analysis.best_checkpoint)
+        ppo_policy = algo.get_policy("learned")
+
+        # Save as torch model
+        ppo_policy.export_model("models")
+        # Save as ONNX model
+        ppo_policy.export_model("models", onnx=11)
+
+        print("Best checkpoint", analysis.best_checkpoint)
 
     ray.shutdown()

connectfour/training/wrappers.py

@@ -1,112 +1,23 @@
 from typing import Optional
 
-from ray.rllib.env.multi_agent_env import MultiAgentEnv
 from ray.rllib.utils.annotations import PublicAPI
-from ray.rllib.utils.gym import convert_old_gym_space_to_gymnasium_space
+from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv
 
 
 @PublicAPI
-class Connect4Env(MultiAgentEnv):
-    """An interface to the PettingZoo MARL environment library.
-
-    See: https://github.com/Farama-Foundation/PettingZoo
-
-    Inherits from MultiAgentEnv and exposes a given AEC
-    (actor-environment-cycle) game from the PettingZoo project via the
-    MultiAgentEnv public API.
-
-    Note that the wrapper has some important limitations:
-
-    1. All agents have the same action_spaces and observation_spaces.
-       Note: If, within your aec game, agents do not have homogeneous action /
-       observation spaces, apply SuperSuit wrappers
-       to apply padding functionality: https://github.com/Farama-Foundation/
-       SuperSuit#built-in-multi-agent-only-functions
-    2. Environments are positive sum games (-> Agents are expected to cooperate
-       to maximize reward). This isn't a hard restriction, it just that
-       standard algorithms aren't expected to work well in highly competitive
-       games."""
-
-    def __init__(self, env):
-        super().__init__()
-        self.env = env
-        env.reset()
-
-        # Since all agents have the same spaces, do not provide full observation-
-        # and action-spaces as Dicts, mapping agent IDs to the individual
-        # agents' spaces. Instead, `self.[action|observation]_space` are the single
-        # agent spaces.
-        self._obs_space_in_preferred_format = False
-        self._action_space_in_preferred_format = False
-
-        # Collect the individual agents' spaces (they should all be the same):
-        first_obs_space = self.env.observation_space(self.env.agents[0])
-        first_action_space = self.env.action_space(self.env.agents[0])
-
-        for agent in self.env.agents:
-            if self.env.observation_space(agent) != first_obs_space:
-                raise ValueError(
-                    "Observation spaces for all agents must be identical. Perhaps "
-                    "SuperSuit's pad_observations wrapper can help (useage: "
-                    "`supersuit.aec_wrappers.pad_observations(env)`"
-                )
-            if self.env.action_space(agent) != first_action_space:
-                raise ValueError(
-                    "Action spaces for all agents must be identical. Perhaps "
-                    "SuperSuit's pad_action_space wrapper can help (usage: "
-                    "`supersuit.aec_wrappers.pad_action_space(env)`"
-                )
-
-        # Convert from gym to gymnasium, if necessary.
-        self.observation_space = convert_old_gym_space_to_gymnasium_space(
-            first_obs_space
-        )
-        self.action_space = convert_old_gym_space_to_gymnasium_space(first_action_space)
-
-        self._agent_ids = set(self.env.agents)
+class Connect4Env(PettingZooEnv):
+    """An interface to the PettingZoo MARL environment library"""
 
     def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
+        # In base class =>
+        # info = self.env.reset(seed=seed, return_info=True, options=options)
         info = self.env.reset(seed=seed, options=options)
         return (
             {self.env.agent_selection: self.env.observe(self.env.agent_selection)},
             info or {},
         )
 
-    def step(self, action):
-        self.env.step(action[self.env.agent_selection])
-        obs_d = {}
-        rew_d = {}
-        terminated_d = {}
-        truncated_d = {}
-        info_d = {}
-        while self.env.agents:
-            obs, rew, terminated, truncated, info = self.env.last()
-            agent_id = self.env.agent_selection
-            obs_d[agent_id] = obs
-            rew_d[agent_id] = rew
-            terminated_d[agent_id] = terminated
-            truncated_d[agent_id] = truncated
-            info_d[agent_id] = info
-            if (
-                self.env.terminations[self.env.agent_selection]
-                or self.env.truncations[self.env.agent_selection]
-            ):
-                self.env.step(None)
-            else:
-                break
-
-        all_gone = not self.env.agents
-        terminated_d["__all__"] = all_gone and all(terminated_d.values())
-        truncated_d["__all__"] = all_gone and all(truncated_d.values())
-
-        return obs_d, rew_d, terminated_d, truncated_d, info_d
-
-    def close(self):
-        self.env.close()
-
     def render(self):
+        # In base class =>
+        # return self.env.render(self.render_mode)
         return self.env.render()
-
-    @property
-    def get_sub_environments(self):
-        return self.env.unwrapped

models/__init__.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+MODEL_PATH = Path(__file__).parent.absolute()

connectfour/checkpoint/algorithm_state.pkl → models/model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cbbc198c3406897931f5f18046a88181b8abff1aedbea1d869329731c9a50853
-size 66321
+oid sha256:ae8bf3e7080eca6ba6c8e68dbceffa03ca8d05f65249c474544a66e039352d2a
+size 361882

poetry.lock

@@ -290,6 +290,18 @@ d = ["aiohttp (>=3.7.4)"]
 jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"]
 uvloop = ["uvloop (>=0.15.2)"]
 
+[[package]]
+name = "cachetools"
+version = "5.3.0"
+description = "Extensible memoizing collections and decorators"
+category = "main"
+optional = false
+python-versions = "~=3.7"
+files = [
+    {file = "cachetools-5.3.0-py3-none-any.whl", hash = "sha256:429e1a1e845c008ea6c85aa35d4b98b65d6a9763eeef3e37e92728a12d1de9d4"},
+    {file = "cachetools-5.3.0.tar.gz", hash = "sha256:13dfddc7b8df938c21a940dfa6557ce6e94a2f1cdfa58eb90c805721d58f2c14"},
+]
+
 [[package]]
 name = "certifi"
 version = "2022.12.7"
@@ -456,6 +468,24 @@ files = [
     {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"},
 ]
 
+[[package]]
+name = "coloredlogs"
+version = "15.0.1"
+description = "Colored terminal output for Python's logging module"
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+files = [
+    {file = "coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934"},
+    {file = "coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0"},
+]
+
+[package.dependencies]
+humanfriendly = ">=9.1"
+
+[package.extras]
+cron = ["capturer (>=2.4)"]
+
 [[package]]
 name = "contourpy"
 version = "1.0.7"
@@ -692,6 +722,18 @@ files = [
 docs = ["furo (>=2022.12.7)", "sphinx (>=6.1.3)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)"]
 testing = ["covdefaults (>=2.3)", "coverage (>=7.2.2)", "diff-cover (>=7.5)", "pytest (>=7.2.2)", "pytest-cov (>=4)", "pytest-mock (>=3.10)", "pytest-timeout (>=2.1)"]
 
+[[package]]
+name = "flatbuffers"
+version = "23.3.3"
+description = "The FlatBuffers serialization format for Python"
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "flatbuffers-23.3.3-py2.py3-none-any.whl", hash = "sha256:5ad36d376240090757e8f0a2cfaf6abcc81c6536c0dc988060375fd0899121f8"},
+    {file = "flatbuffers-23.3.3.tar.gz", hash = "sha256:cabd87c4882f37840f6081f094b2c5bc28cefc2a6357732746936d055ab45c3d"},
+]
+
 [[package]]
 name = "fonttools"
 version = "4.39.3"
@@ -848,16 +890,60 @@ files = [
     {file = "gast-0.5.3.tar.gz", hash = "sha256:cfbea25820e653af9c7d1807f659ce0a0a9c64f2439421a7bba4f0983f532dea"},
 ]
 
+[[package]]
+name = "google-auth"
+version = "2.17.1"
+description = "Google Authentication Library"
+category = "main"
+optional = false
+python-versions = ">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*"
+files = [
+    {file = "google-auth-2.17.1.tar.gz", hash = "sha256:8f379b46bad381ad2a0b989dfb0c13ad28d3c2a79f27348213f8946a1d15d55a"},
+    {file = "google_auth-2.17.1-py2.py3-none-any.whl", hash = "sha256:357ff22a75b4c0f6093470f21816a825d2adee398177569824e37b6c10069e19"},
+]
+
+[package.dependencies]
+cachetools = ">=2.0.0,<6.0"
+pyasn1-modules = ">=0.2.1"
+rsa = {version = ">=3.1.4,<5", markers = "python_version >= \"3.6\""}
+six = ">=1.9.0"
+
+[package.extras]
+aiohttp = ["aiohttp (>=3.6.2,<4.0.0dev)", "requests (>=2.20.0,<3.0.0dev)"]
+enterprise-cert = ["cryptography (==36.0.2)", "pyopenssl (==22.0.0)"]
+pyopenssl = ["cryptography (>=38.0.3)", "pyopenssl (>=20.0.0)"]
+reauth = ["pyu2f (>=0.1.5)"]
+requests = ["requests (>=2.20.0,<3.0.0dev)"]
+
+[[package]]
+name = "google-auth-oauthlib"
+version = "0.4.6"
+description = "Google Authentication Library"
+category = "main"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "google-auth-oauthlib-0.4.6.tar.gz", hash = "sha256:a90a072f6993f2c327067bf65270046384cda5a8ecb20b94ea9a687f1f233a7a"},
+    {file = "google_auth_oauthlib-0.4.6-py2.py3-none-any.whl", hash = "sha256:3f2a6e802eebbb6fb736a370fbf3b055edcb6b52878bf2f26330b5e041316c73"},
+]
+
+[package.dependencies]
+google-auth = ">=1.0.0"
+requests-oauthlib = ">=0.7.0"
+
+[package.extras]
+tool = ["click (>=6.0.0)"]
+
 [[package]]
 name = "gradio"
-version = "3.23.0"
+version = "3.24.0"
 description = "Python library for easily interacting with trained machine learning models"
 category = "main"
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "gradio-3.23.0-py3-none-any.whl", hash = "sha256:1f637f80e4b740892c3cc5ec559158afe2b6b97f0a18eca588c235ddedddc7ec"},
-    {file = "gradio-3.23.0.tar.gz", hash = "sha256:bfb9f59d799271029e6309207c7f42ed8adb297dcb8bbe31c83c3ebf32dbe193"},
+    {file = "gradio-3.24.0-py3-none-any.whl", hash = "sha256:cedd67f7cbd17764b3613fb4df274a7c450c74e31a2e3229097d43cb4ffa50c7"},
+    {file = "gradio-3.24.0.tar.gz", hash = "sha256:4ac2bf531b3c0ff5ec9e93959f2d1dbc49eac1767bafa2d80f8950a3bc40c4ed"},
 ]
 
 [package.dependencies]
@@ -866,7 +952,7 @@ aiohttp = "*"
 altair = ">=4.2.0"
 fastapi = "*"
 ffmpy = "*"
-fsspec = "*"
+gradio-client = ">=0.0.5"
 httpx = "*"
 huggingface-hub = ">=0.13.0"
 jinja2 = "*"
@@ -888,6 +974,25 @@ typing-extensions = "*"
 uvicorn = "*"
 websockets = ">=10.0"
 
+[[package]]
+name = "gradio-client"
+version = "0.0.5"
+description = "Python library for easily interacting with trained machine learning models"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "gradio_client-0.0.5-py3-none-any.whl", hash = "sha256:ca4167ebae72d920ebec2be47010cf60e31e0296ad9baac771befb17b87f0eef"},
+    {file = "gradio_client-0.0.5.tar.gz", hash = "sha256:dc6479a119314aac0bbf6821da6e946df17f048cc571559379a89590618f7b5d"},
+]
+
+[package.dependencies]
+fsspec = "*"
+huggingface-hub = ">=0.13.0"
+packaging = "*"
+requests = "*"
+websockets = "*"
+
 [[package]]
 name = "grpcio"
 version = "1.49.1"
@@ -1138,6 +1243,21 @@ testing = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "jedi", "pytest", "pytest
 torch = ["torch"]
 typing = ["types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3"]
 
+[[package]]
+name = "humanfriendly"
+version = "10.0"
+description = "Human friendly output for text interfaces using Python"
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+files = [
+    {file = "humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477"},
+    {file = "humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc"},
+]
+
+[package.dependencies]
+pyreadline3 = {version = "*", markers = "sys_platform == \"win32\" and python_version >= \"3.8\""}
+
 [[package]]
 name = "idna"
 version = "3.4"
@@ -1531,6 +1651,24 @@ docs = ["sphinx (>=1.6.0)", "sphinx-bootstrap-theme"]
 flake8 = ["flake8"]
 tests = ["psutil", "pytest (!=3.3.0)", "pytest-cov"]
 
+[[package]]
+name = "markdown"
+version = "3.4.3"
+description = "Python implementation of John Gruber's Markdown."
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "Markdown-3.4.3-py3-none-any.whl", hash = "sha256:065fd4df22da73a625f14890dd77eb8040edcbd68794bcd35943be14490608b2"},
+    {file = "Markdown-3.4.3.tar.gz", hash = "sha256:8bf101198e004dc93e84a12a7395e31aac6a9c9942848ae1d99b9d72cf9b3520"},
+]
+
+[package.dependencies]
+importlib-metadata = {version = ">=4.4", markers = "python_version < \"3.10\""}
+
+[package.extras]
+testing = ["coverage", "pyyaml"]
+
 [[package]]
 name = "markdown-it-py"
 version = "2.2.0"
@@ -2184,6 +2322,111 @@ files = [
 setuptools = "*"
 wheel = "*"
 
+[[package]]
+name = "oauthlib"
+version = "3.2.2"
+description = "A generic, spec-compliant, thorough implementation of the OAuth request-signing logic"
+category = "main"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "oauthlib-3.2.2-py3-none-any.whl", hash = "sha256:8139f29aac13e25d502680e9e19963e83f16838d48a0d71c287fe40e7067fbca"},
+    {file = "oauthlib-3.2.2.tar.gz", hash = "sha256:9859c40929662bec5d64f34d01c99e093149682a3f38915dc0655d5a633dd918"},
+]
+
+[package.extras]
+rsa = ["cryptography (>=3.0.0)"]
+signals = ["blinker (>=1.4.0)"]
+signedtoken = ["cryptography (>=3.0.0)", "pyjwt (>=2.0.0,<3)"]
+
+[[package]]
+name = "onnx"
+version = "1.12.0"
+description = "Open Neural Network Exchange"
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "onnx-1.12.0-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:bdbd2578424c70836f4d0f9dda16c21868ddb07cc8192f9e8a176908b43d694b"},
+    {file = "onnx-1.12.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:213e73610173f6b2e99f99a4b0636f80b379c417312079d603806e48ada4ca8b"},
+    {file = "onnx-1.12.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9fd2f4e23078df197bb76a59b9cd8f5a43a6ad2edc035edb3ecfb9042093e05a"},
+    {file = "onnx-1.12.0-cp310-cp310-win32.whl", hash = "sha256:23781594bb8b7ee985de1005b3c601648d5b0568a81e01365c48f91d1f5648e4"},
+    {file = "onnx-1.12.0-cp310-cp310-win_amd64.whl", hash = "sha256:81a3555fd67be2518bf86096299b48fb9154652596219890abfe90bd43a9ec13"},
+    {file = "onnx-1.12.0-cp37-cp37m-macosx_10_12_x86_64.whl", hash = "sha256:5578b93dc6c918cec4dee7fb7d9dd3b09d338301ee64ca8b4f28bc217ed42dca"},
+    {file = "onnx-1.12.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c11162ffc487167da140f1112f49c4f82d815824f06e58bc3095407699f05863"},
+    {file = "onnx-1.12.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:341c7016e23273e9ffa9b6e301eee95b8c37d0f04df7cedbdb169d2c39524c96"},
+    {file = "onnx-1.12.0-cp37-cp37m-win32.whl", hash = "sha256:3c6e6bcffc3f5c1e148df3837dc667fa4c51999788c1b76b0b8fbba607e02da8"},
+    {file = "onnx-1.12.0-cp37-cp37m-win_amd64.whl", hash = "sha256:8a7aa61aea339bd28f310f4af4f52ce6c4b876386228760b16308efd58f95059"},
+    {file = "onnx-1.12.0-cp38-cp38-macosx_10_12_x86_64.whl", hash = "sha256:56ceb7e094c43882b723cfaa107d85ad673cfdf91faeb28d7dcadacca4f43a07"},
+    {file = "onnx-1.12.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b3629e8258db15d4e2c9b7f1be91a3186719dd94661c218c6f5fde3cc7de3d4d"},
+    {file = "onnx-1.12.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2d9a7db54e75529160337232282a4816cc50667dc7dc34be178fd6f6b79d4705"},
+    {file = "onnx-1.12.0-cp38-cp38-win32.whl", hash = "sha256:fea5156a03398fe0e23248042d8651c1eaac5f6637d4dd683b4c1f1320b9f7b4"},
+    {file = "onnx-1.12.0-cp38-cp38-win_amd64.whl", hash = "sha256:f66d2996e65f490a57b3ae952e4e9189b53cc9fe3f75e601d50d4db2dc1b1cd9"},
+    {file = "onnx-1.12.0-cp39-cp39-macosx_10_12_x86_64.whl", hash = "sha256:c39a7a0352c856f1df30dccf527eb6cb4909052e5eaf6fa2772a637324c526aa"},
+    {file = "onnx-1.12.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fab13feb4d94342aae6d357d480f2e47d41b9f4e584367542b21ca6defda9e0a"},
+    {file = "onnx-1.12.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c7a9b3ea02c30efc1d2662337e280266aca491a8e86be0d8a657f874b7cccd1e"},
+    {file = "onnx-1.12.0-cp39-cp39-win32.whl", hash = "sha256:f8800f28c746ab06e51ef8449fd1215621f4ddba91be3ffc264658937d38a2af"},
+    {file = "onnx-1.12.0-cp39-cp39-win_amd64.whl", hash = "sha256:af90427ca04c6b7b8107c2021e1273227a3ef1a7a01f3073039cae7855a59833"},
+    {file = "onnx-1.12.0.tar.gz", hash = "sha256:13b3e77d27523b9dbf4f30dfc9c959455859d5e34e921c44f712d69b8369eff9"},
+]
+
+[package.dependencies]
+numpy = ">=1.16.6"
+protobuf = ">=3.12.2,<=3.20.1"
+typing-extensions = ">=3.6.2.1"
+
+[package.extras]
+lint = ["clang-format (==13.0.0)", "flake8", "mypy (==0.782)", "types-protobuf (==3.18.4)"]
+
+[[package]]
+name = "onnxruntime"
+version = "1.14.1"
+description = "ONNX Runtime is a runtime accelerator for Machine Learning models"
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "onnxruntime-1.14.1-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:193ef1ac512e530c6e6e259c26e67212e2cd3f2bfaad6ff935ed3f4281053056"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d2853bbb36cb272d99f6c225e5040eb0ddb37a667fce20d186ecdf0a6fac8af8"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8e1b173365c6894616b8207e23cbb891da9638c5373668d6653e4081ef5f04d0"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:24bf0401c5f92be7230ac660ff07ba06f7c175e99e225d5d48ff09062a3b76e9"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-manylinux_2_27_aarch64.whl", hash = "sha256:0a2d09260bbdbe1df678e0a237a5f7b1a44fd11a2f52688d8b6a53a9d03a26db"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-manylinux_2_27_x86_64.whl", hash = "sha256:d99d35b9d5c3f46cad1673a39cc753fb57d60784369b59e6f8cd3dfb77df1885"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-win32.whl", hash = "sha256:f400356df1b27d9adc5513319e8a89753e48ef0d6c5084caf5db8e132f46e7e8"},
+    {file = "onnxruntime-1.14.1-cp310-cp310-win_amd64.whl", hash = "sha256:96a4059dbab162fe5cdb6750f8c70b2106ef2de5d49a7f72085171937d0e36d3"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-macosx_10_15_x86_64.whl", hash = "sha256:fa23df6a349218636290f9fe56d7baaceb1a50cf92255234d495198b47d92327"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bc70e44d9e123d126648da24ffb39e56464272a1660a3eb91f4f5b74263be3ba"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:deff8138045a3affb6be064b598e3ec69a88e4d445359c50464ee5379b8eaf19"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-manylinux_2_27_aarch64.whl", hash = "sha256:7c02acdc1107cbf698dcbf6dadc6f5b6aa179e7fa9a026251e99cf8613bd3129"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-manylinux_2_27_x86_64.whl", hash = "sha256:6efa3b2f4b1eaa6c714c07861993bfd9bb33bd73cdbcaf5b4aadcf1ec13fcaf7"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-win32.whl", hash = "sha256:72fc0acc82c54bf03eba065ad9025baa438c00c54a2ee0beb8ae4b6085cd3a0d"},
+    {file = "onnxruntime-1.14.1-cp37-cp37m-win_amd64.whl", hash = "sha256:4d6f08ea40d63ccf90f203f4a2a498f4e590737dcaf16867075cc8e0a86c5554"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:c2d9e8f1bc6037f14d8aaa480492792c262fc914936153e40b06b3667bb25549"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:e7424d3befdd95b537c90787bbfaa053b2bb19eb60135abb898cb0e099d7d7ad"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9066d275e6e41d0597e234d2d88c074d4325e650c74a9527a52cadbcf42a0fe2"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8224d3c1f2cd0b899cea7b5a39f28b971debe0da30fcbc61382801d97d6f5740"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-manylinux_2_27_aarch64.whl", hash = "sha256:f4ac52ff4ac793683ebd1fbd1ee24197e3b4ca825ee68ff739296a820867debe"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-manylinux_2_27_x86_64.whl", hash = "sha256:b1dd8cdd3be36c32ddd8f5763841ed571c3e81da59439a622947bd97efee6e77"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-win32.whl", hash = "sha256:95d0f0cd95360c07f1c3ba20962b9bb813627df4bfc1b4b274e1d40044df5ad1"},
+    {file = "onnxruntime-1.14.1-cp38-cp38-win_amd64.whl", hash = "sha256:de40a558e00fc00f92e298d5be99eb8075dba51368dabcb259670a00f4670e56"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-macosx_10_15_x86_64.whl", hash = "sha256:c65b587a42a89fceceaad367bd69d071ee5c9c7010b76e2adac5e9efd9356fb5"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:6e47ef6a2c6e6dd6ff48bc13f2331d124dff00e1d76627624bb3268c8058f19c"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0afd0f671d068dd99b9d071d88e93a9a57a5ed59af440c0f4d65319ee791603f"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc65e9061349cdf98ce16b37722b557109f16076632fbfed9a3151895cfd3bb7"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-manylinux_2_27_aarch64.whl", hash = "sha256:2ff17c71187391a71e6ccc78ca89aed83bcaed1c085c95267ab1a70897868bdd"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-manylinux_2_27_x86_64.whl", hash = "sha256:9b795189916942ce848192200dde5b1f32799ee6c84fc600969a44d88e8a5404"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-win32.whl", hash = "sha256:17ca3100112af045118750d24643a01ed4e6d86071a8efaef75cc1d434ea64aa"},
+    {file = "onnxruntime-1.14.1-cp39-cp39-win_amd64.whl", hash = "sha256:b5e8c489329ba0fa0639dfd7ec02d6b07cece1bab52ef83884b537247efbda74"},
+]
+
+[package.dependencies]
+coloredlogs = "*"
+flatbuffers = "*"
+numpy = ">=1.21.6"
+packaging = "*"
+protobuf = "*"
+sympy = "*"
+
 [[package]]
 name = "orjson"
 version = "3.8.8"
@@ -2471,39 +2714,65 @@ testing = ["pytest", "pytest-benchmark"]
 
 [[package]]
 name = "protobuf"
-version = "3.17.0"
+version = "3.19.6"
 description = "Protocol Buffers"
 category = "main"
 optional = false
+python-versions = ">=3.5"
+files = [
+    {file = "protobuf-3.19.6-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:010be24d5a44be7b0613750ab40bc8b8cedc796db468eae6c779b395f50d1fa1"},
+    {file = "protobuf-3.19.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:11478547958c2dfea921920617eb457bc26867b0d1aa065ab05f35080c5d9eb6"},
+    {file = "protobuf-3.19.6-cp310-cp310-win32.whl", hash = "sha256:559670e006e3173308c9254d63facb2c03865818f22204037ab76f7a0ff70b5f"},
+    {file = "protobuf-3.19.6-cp310-cp310-win_amd64.whl", hash = "sha256:347b393d4dd06fb93a77620781e11c058b3b0a5289262f094379ada2920a3730"},
+    {file = "protobuf-3.19.6-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:a8ce5ae0de28b51dff886fb922012dad885e66176663950cb2344c0439ecb473"},
+    {file = "protobuf-3.19.6-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90b0d02163c4e67279ddb6dc25e063db0130fc299aefabb5d481053509fae5c8"},
+    {file = "protobuf-3.19.6-cp36-cp36m-win32.whl", hash = "sha256:30f5370d50295b246eaa0296533403961f7e64b03ea12265d6dfce3a391d8992"},
+    {file = "protobuf-3.19.6-cp36-cp36m-win_amd64.whl", hash = "sha256:0c0714b025ec057b5a7600cb66ce7c693815f897cfda6d6efb58201c472e3437"},
+    {file = "protobuf-3.19.6-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:5057c64052a1f1dd7d4450e9aac25af6bf36cfbfb3a1cd89d16393a036c49157"},
+    {file = "protobuf-3.19.6-cp37-cp37m-manylinux2014_aarch64.whl", hash = "sha256:bb6776bd18f01ffe9920e78e03a8676530a5d6c5911934c6a1ac6eb78973ecb6"},
+    {file = "protobuf-3.19.6-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:84a04134866861b11556a82dd91ea6daf1f4925746b992f277b84013a7cc1229"},
+    {file = "protobuf-3.19.6-cp37-cp37m-win32.whl", hash = "sha256:4bc98de3cdccfb5cd769620d5785b92c662b6bfad03a202b83799b6ed3fa1fa7"},
+    {file = "protobuf-3.19.6-cp37-cp37m-win_amd64.whl", hash = "sha256:aa3b82ca1f24ab5326dcf4ea00fcbda703e986b22f3d27541654f749564d778b"},
+    {file = "protobuf-3.19.6-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:2b2d2913bcda0e0ec9a784d194bc490f5dc3d9d71d322d070b11a0ade32ff6ba"},
+    {file = "protobuf-3.19.6-cp38-cp38-manylinux2014_aarch64.whl", hash = "sha256:d0b635cefebd7a8a0f92020562dead912f81f401af7e71f16bf9506ff3bdbb38"},
+    {file = "protobuf-3.19.6-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a552af4dc34793803f4e735aabe97ffc45962dfd3a237bdde242bff5a3de684"},
+    {file = "protobuf-3.19.6-cp38-cp38-win32.whl", hash = "sha256:0469bc66160180165e4e29de7f445e57a34ab68f49357392c5b2f54c656ab25e"},
+    {file = "protobuf-3.19.6-cp38-cp38-win_amd64.whl", hash = "sha256:91d5f1e139ff92c37e0ff07f391101df77e55ebb97f46bbc1535298d72019462"},
+    {file = "protobuf-3.19.6-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:c0ccd3f940fe7f3b35a261b1dd1b4fc850c8fde9f74207015431f174be5976b3"},
+    {file = "protobuf-3.19.6-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:30a15015d86b9c3b8d6bf78d5b8c7749f2512c29f168ca259c9d7727604d0e39"},
+    {file = "protobuf-3.19.6-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:878b4cd080a21ddda6ac6d1e163403ec6eea2e206cf225982ae04567d39be7b0"},
+    {file = "protobuf-3.19.6-cp39-cp39-win32.whl", hash = "sha256:5a0d7539a1b1fb7e76bf5faa0b44b30f812758e989e59c40f77a7dab320e79b9"},
+    {file = "protobuf-3.19.6-cp39-cp39-win_amd64.whl", hash = "sha256:bbf5cea5048272e1c60d235c7bd12ce1b14b8a16e76917f371c718bd3005f045"},
+    {file = "protobuf-3.19.6-py2.py3-none-any.whl", hash = "sha256:14082457dc02be946f60b15aad35e9f5c69e738f80ebbc0900a19bc83734a5a4"},
+    {file = "protobuf-3.19.6.tar.gz", hash = "sha256:5f5540d57a43042389e87661c6eaa50f47c19c6176e8cf1c4f287aeefeccb5c4"},
+]
+
+[[package]]
+name = "pyasn1"
+version = "0.4.8"
+description = "ASN.1 types and codecs"
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "pyasn1-0.4.8-py2.py3-none-any.whl", hash = "sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d"},
+    {file = "pyasn1-0.4.8.tar.gz", hash = "sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba"},
+]
+
+[[package]]
+name = "pyasn1-modules"
+version = "0.2.8"
+description = "A collection of ASN.1-based protocols modules."
+category = "main"
+optional = false
 python-versions = "*"
 files = [
-    {file = "protobuf-3.17.0-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:15351df904347da2081a2eebc42b192c29724eb57dbe56dae440be843f1e4779"},
-    {file = "protobuf-3.17.0-cp27-cp27mu-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:5356981c1919782b8c2e3ea5c5d85ad5937b8178a025ac9edc2f2ca5b4a717ae"},
-    {file = "protobuf-3.17.0-cp35-cp35m-macosx_10_9_intel.whl", hash = "sha256:eac0a2a7ea99e17175f6e7b53cdc9004ed786c072fbdf933def0e454e14fd323"},
-    {file = "protobuf-3.17.0-cp35-cp35m-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:4c8d0997fdc0a4cf9de7950d598ce6974b22e8618bbcf1d15e9842010cf8420a"},
-    {file = "protobuf-3.17.0-cp35-cp35m-win32.whl", hash = "sha256:9ae321459d4890c3939c536382f75e232c9e91ce506310353c8a15ad5c379e0d"},
-    {file = "protobuf-3.17.0-cp35-cp35m-win_amd64.whl", hash = "sha256:295944ef0772498d7bf75f6aa5d4dfcfd02f5ce70f735b406e52e43ac3914d38"},
-    {file = "protobuf-3.17.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:850f429bd2399525d339d05bc809f090f16d3d88737bed637d355a5ee8d3b81a"},
-    {file = "protobuf-3.17.0-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:809a96d5a1a74538728710f9104f43ae77f5e48bde274ee321b10a324ba52e4f"},
-    {file = "protobuf-3.17.0-cp36-cp36m-win32.whl", hash = "sha256:8a3ac375539055164f31a330770f137875307e6f04c21e2647f2e7139c501295"},
-    {file = "protobuf-3.17.0-cp36-cp36m-win_amd64.whl", hash = "sha256:3d338910b10b88b18581cf6877b3938b2e262e8fdc2c1057f5a291787de63183"},
-    {file = "protobuf-3.17.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:1488f786bd1912f97796cf5def8cacf433735616896cf7ed9dc786cee693dfc8"},
-    {file = "protobuf-3.17.0-cp37-cp37m-manylinux2014_aarch64.whl", hash = "sha256:bcaff977db178f0bfde10bab0d23a5f5adf5964adba70c315e45922a1c55eb90"},
-    {file = "protobuf-3.17.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:939ce06846ddfec99c0bff510510b3ee45778e7a3aec6544d1f36526e5fecb67"},
-    {file = "protobuf-3.17.0-cp37-cp37m-win32.whl", hash = "sha256:3237acce5b666c7b0f45785cc2d0809796d4df3593bd68338aebf25408139188"},
-    {file = "protobuf-3.17.0-cp37-cp37m-win_amd64.whl", hash = "sha256:2f77afe33bb86c7d34221a86193256d69aa10818620fe4a7513d98211d67d672"},
-    {file = "protobuf-3.17.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:acc9f2091ace3de429eee424ab7ba0bc52a6aa9ffc9909e5c4de259a3f71db46"},
-    {file = "protobuf-3.17.0-cp38-cp38-manylinux2014_aarch64.whl", hash = "sha256:a29631f4f8bcf79b12a59e83d238d888de5034871461d788c74c68218ad75049"},
-    {file = "protobuf-3.17.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:05c304396e309661c45e3a97bd2d8da1fc2bab743ed2ca880bcb757271c40c0e"},
-    {file = "protobuf-3.17.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:baea44967071e6a51e705e4e88aebf35f530a14004cc69f60a185e5d7e13de7e"},
-    {file = "protobuf-3.17.0-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:3b5c461af5a3cebd796c73370db929b7e24cbaba655eefdc044226bc8a843d6b"},
-    {file = "protobuf-3.17.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:44399393c3a8cc04a4cfbdc721dd7f2114497efda582e946a91b8c4290ae5ff5"},
-    {file = "protobuf-3.17.0-py2.py3-none-any.whl", hash = "sha256:e32ef0c9f4b548c80d94dfff8b4130ca2ff3d50caaf2455889e3f5b8a01e8038"},
-    {file = "protobuf-3.17.0.tar.gz", hash = "sha256:05dfe9319939a8473c21b469f34f6486646e54fb8542637cf7ed8e2fbfe21538"},
-]
-
-[package.dependencies]
-six = ">=1.9"
+    {file = "pyasn1-modules-0.2.8.tar.gz", hash = "sha256:905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e"},
+    {file = "pyasn1_modules-0.2.8-py2.py3-none-any.whl", hash = "sha256:a50b808ffeb97cb3601dd25981f6b016cbb3d31fbf57a8b8a87428e6158d0c74"},
+]
+
+[package.dependencies]
+pyasn1 = ">=0.4.6,<0.5.0"
 
 [[package]]
 name = "pydantic"
@@ -2705,6 +2974,18 @@ files = [
 [package.extras]
 diagrams = ["jinja2", "railroad-diagrams"]
 
+[[package]]
+name = "pyreadline3"
+version = "3.4.1"
+description = "A python implementation of GNU readline."
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "pyreadline3-3.4.1-py3-none-any.whl", hash = "sha256:b0efb6516fd4fb07b45949053826a62fa4cb353db5be2bbb4a7aa1fdd1e345fb"},
+    {file = "pyreadline3-3.4.1.tar.gz", hash = "sha256:6f3d1f7b8a31ba32b73917cefc1f28cc660562f39aea8646d30bd6eff21f7bae"},
+]
+
 [[package]]
 name = "pyrsistent"
 version = "0.19.3"
@@ -2997,6 +3278,25 @@ urllib3 = ">=1.21.1,<1.27"
 socks = ["PySocks (>=1.5.6,!=1.5.7)"]
 use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
 
+[[package]]
+name = "requests-oauthlib"
+version = "1.3.1"
+description = "OAuthlib authentication support for Requests."
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
+files = [
+    {file = "requests-oauthlib-1.3.1.tar.gz", hash = "sha256:75beac4a47881eeb94d5ea5d6ad31ef88856affe2332b9aafb52c6452ccf0d7a"},
+    {file = "requests_oauthlib-1.3.1-py2.py3-none-any.whl", hash = "sha256:2577c501a2fb8d05a304c09d090d6e47c306fef15809d102b327cf8364bddab5"},
+]
+
+[package.dependencies]
+oauthlib = ">=3.0.0"
+requests = ">=2.0.0"
+
+[package.extras]
+rsa = ["oauthlib[signedtoken] (>=3.0.0)"]
+
 [[package]]
 name = "rfc3986"
 version = "1.5.0"
@@ -3035,6 +3335,21 @@ typing-extensions = {version = ">=4.0.0,<5.0", markers = "python_version < \"3.9
 [package.extras]
 jupyter = ["ipywidgets (>=7.5.1,<9)"]
 
+[[package]]
+name = "rsa"
+version = "4.9"
+description = "Pure-Python RSA implementation"
+category = "main"
+optional = false
+python-versions = ">=3.6,<4"
+files = [
+    {file = "rsa-4.9-py3-none-any.whl", hash = "sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7"},
+    {file = "rsa-4.9.tar.gz", hash = "sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21"},
+]
+
+[package.dependencies]
+pyasn1 = ">=0.1.3"
+
 [[package]]
 name = "scikit-image"
 version = "0.20.0"
@@ -3231,6 +3546,56 @@ files = [
 [package.extras]
 widechars = ["wcwidth"]
 
+[[package]]
+name = "tensorboard"
+version = "2.12.0"
+description = "TensorBoard lets you watch Tensors Flow"
+category = "main"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "tensorboard-2.12.0-py3-none-any.whl", hash = "sha256:3cbdc32448d7a28dc1bf0b1754760c08b8e0e2e37c451027ebd5ff4896613012"},
+]
+
+[package.dependencies]
+absl-py = ">=0.4"
+google-auth = ">=1.6.3,<3"
+google-auth-oauthlib = ">=0.4.1,<0.5"
+grpcio = ">=1.48.2"
+markdown = ">=2.6.8"
+numpy = ">=1.12.0"
+protobuf = ">=3.19.6"
+requests = ">=2.21.0,<3"
+setuptools = ">=41.0.0"
+tensorboard-data-server = ">=0.7.0,<0.8.0"
+tensorboard-plugin-wit = ">=1.6.0"
+werkzeug = ">=1.0.1"
+wheel = ">=0.26"
+
+[[package]]
+name = "tensorboard-data-server"
+version = "0.7.0"
+description = "Fast data loading for TensorBoard"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "tensorboard_data_server-0.7.0-py3-none-any.whl", hash = "sha256:753d4214799b31da7b6d93837959abebbc6afa86e69eacf1e9a317a48daa31eb"},
+    {file = "tensorboard_data_server-0.7.0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:eb7fa518737944dbf4f0cf83c2e40a7ac346bf91be2e6a0215de98be74e85454"},
+    {file = "tensorboard_data_server-0.7.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64aa1be7c23e80b1a42c13b686eb0875bb70f5e755f4d2b8de5c1d880cf2267f"},
+]
+
+[[package]]
+name = "tensorboard-plugin-wit"
+version = "1.8.1"
+description = "What-If Tool TensorBoard plugin."
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "tensorboard_plugin_wit-1.8.1-py3-none-any.whl", hash = "sha256:ff26bdd583d155aa951ee3b152b3d0cffae8005dc697f72b44a8e8c2a77a8cbe"},
+]
+
 [[package]]
 name = "tensorboardx"
 version = "2.6"
@@ -3625,6 +3990,24 @@ files = [
     {file = "websockets-10.4.tar.gz", hash = "sha256:eef610b23933c54d5d921c92578ae5f89813438fded840c2e9809d378dc765d3"},
 ]
 
+[[package]]
+name = "werkzeug"
+version = "2.2.3"
+description = "The comprehensive WSGI web application library."
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "Werkzeug-2.2.3-py3-none-any.whl", hash = "sha256:56433961bc1f12533306c624f3be5e744389ac61d722175d543e1751285da612"},
+    {file = "Werkzeug-2.2.3.tar.gz", hash = "sha256:2e1ccc9417d4da358b9de6f174e3ac094391ea1d4fbef2d667865d819dfd0afe"},
+]
+
+[package.dependencies]
+MarkupSafe = ">=2.1.1"
+
+[package.extras]
+watchdog = ["watchdog"]
+
 [[package]]
 name = "wheel"
 version = "0.40.0"
@@ -3832,4 +4215,4 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8,<3.11"
-content-hash = "6ac937208f9e895063d08a5437db35ed4c5d75275bcf70e8503c609f42b8f302"
+content-hash = "81eac0c68b289dd9d22d17dd34ae6bcd31f53c8201f229d31ff96d7094ad2392"

pyproject.toml

@@ -23,8 +23,11 @@ pygame = "^2.3.0"
 torch = "^2.0.0"
 libclang = "15.0.6.1"
 tensorflow-probability = "^0.19.0"
-protobuf = "3.17.0"
+protobuf = "3.19.6"
 scipy = ">=1.8,<1.9.2"
+onnx = "1.12.0"
+tensorboard = "^2.12.0"
+onnxruntime = "^1.14.1"
 
 [tool.poetry.dev-dependencies]
 pylint = "*"