AlphaZero and Ranked reward implementation (#6385)

ci/jenkins_tests/run_rllib_tests.sh

@@ -426,6 +426,9 @@ docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
 docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
     /ray/ci/suppress_output --force-direct python /ray/rllib/contrib/random_agent/random_agent.py
 
+docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
+    /ray/ci/suppress_output --force-direct python /ray/rllib/contrib/alpha_zero/examples/train_cartpole.py --training-iteration=1
+
 docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
     /ray/ci/suppress_output --force-direct python /ray/rllib/examples/centralized_critic.py --stop=2000
 

doc/source/rllib-algorithms.rst

@@ -364,3 +364,17 @@ Tuned examples: `CartPole-v0 <https://github.com/ray-project/ray/blob/master/rll
    :language: python
    :start-after: __sphinx_doc_begin__
    :end-before: __sphinx_doc_end__
+
+Single-Player Alpha Zero (contrib/AlphaZero)
+--------------------------------------------
+
+`[paper] <https://arxiv.org/abs/1712.01815>`__ `[implementation] <https://github.com/ray-project/ray/blob/master/rllib/contrib/alpha_zero>`__ AlphaZero is an RL agent originally designed for two-player games. This version adapts it to handle single player games. The code can be used with the SyncSamplesOptimizer as well as with a modified version of the SyncReplayOptimizer, and it scales to any number of workers. It also implements the ranked rewards `(R2) <https://arxiv.org/abs/1807.01672>`__ strategy to enable self-play even in the one-player setting. The code is mainly purposed to be used for combinatorial optimization.
+
+Tuned examples: `CartPole-v0 <https://github.com/ray-project/ray/blob/master/rllib/contrib/alpha_zero/examples/train_cartpole.py>`__
+
+**AlphaZero-specific configs** (see also `common configs <rllib-training.html#common-parameters>`__):
+
+.. literalinclude:: ../../rllib/contrib/alpha_zero/core/alpha_zero_trainer.py
+   :language: python
+   :start-after: __sphinx_doc_begin__
+   :end-before: __sphinx_doc_end__

doc/source/rllib-toc.rst

@@ -78,6 +78,8 @@ Algorithms
 
    -  `Asynchronous Proximal Policy Optimization (APPO) <rllib-algorithms.html#asynchronous-proximal-policy-optimization-appo>`__
 
+   -  `Single-Player AlphaZero (contrib/AlphaZero) <rllib-algorithms.html#single-player-alpha-zero-contrib-alphazero>`__
+
 *  Gradient-based
 
    -  `Advantage Actor-Critic (A2C, A3C) <rllib-algorithms.html#advantage-actor-critic-a2c-a3c>`__

rllib/contrib/alpha_zero/README.md

@@ -0,0 +1,24 @@
+# AlphaZero implementation for Ray/RLlib
+## Notes
+
+This code implements a one-player AlphaZero agent. It includes the "ranked rewards" (R2) strategy which simulates the self-play in the two-player AlphaZero in forcing the agent to be better than its previous self. R2 is also very helpful to normalize dynamically the rewards. 
+
+The code is Pytorch based. It assumes that the environment is a gym environment, has a discrete action space and returns an observation as a dictionary with two keys:
+
+ - `obs` that contains an observation under either the form of a state vector or an image
+ - `action_mask` that contains a mask over the legal actions
+ 
+ It should also implement a `get_state`and a `set_state` function.
+ 
+ The model used in AlphaZero trainer should extend `ActorCriticModel` and implement the method `compute_priors_and_value`. 
+ 
+## Example on Cartpole
+
+Note that both mean and max rewards are obtained with the MCTS in exploration mode: dirichlet noise is added to priors and actions are sampled from the tree policy vectors. We will add later the display of the MCTS in exploitation mode: no dirichlet noise and actions are chosen as tree policy vectors argmax.
+![cartpole_plot](doc/cartpole_plot.png)
+
+## References
+
+- AlphaZero: https://arxiv.org/abs/1712.01815
+- Ranked rewards: https://arxiv.org/abs/1807.01672
+    

rllib/contrib/alpha_zero/core/alpha_zero_policy.py

@@ -0,0 +1,126 @@
+import numpy as np
+import torch
+from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY
+from ray.rllib.policy.torch_policy import TorchPolicy
+from ray.rllib.utils.annotations import override
+
+from ray.rllib.contrib.alpha_zero.core.mcts import Node, RootParentNode
+
+
+class AlphaZeroPolicy(TorchPolicy):
+    def __init__(self, observation_space, action_space, model, loss,
+                 action_distribution_class, mcts_creator, env_creator,
+                 **kwargs):
+        super().__init__(observation_space, action_space, model, loss,
+                         action_distribution_class)
+        # we maintain an env copy in the policy that is used during mcts
+        # simulations
+        self.env_creator = env_creator
+        self.mcts = mcts_creator()
+        self.env = self.env_creator()
+        self.env.reset()
+        self.obs_space = observation_space
+
+    @override(TorchPolicy)
+    def compute_actions(self,
+                        obs_batch,
+                        state_batches=None,
+                        prev_action_batch=None,
+                        prev_reward_batch=None,
+                        info_batch=None,
+                        episodes=None,
+                        **kwargs):
+
+        with torch.no_grad():
+            input_dict = {"obs": obs_batch}
+            if prev_action_batch:
+                input_dict["prev_actions"] = prev_action_batch
+            if prev_reward_batch:
+                input_dict["prev_rewards"] = prev_reward_batch
+
+            actions = []
+
+            for i, episode in enumerate(episodes):
+                if episode.length == 0:
+                    # if first time step of episode, get initial env state
+                    env_state = episode.user_data["initial_state"]
+                    # verify if env has been wrapped for ranked rewards
+                    if self.env.__class__.__name__ == \
+                            "RankedRewardsEnvWrapper":
+                        # r2 env state contains also the rewards buffer state
+                        env_state = {
+                            "env_state": env_state,
+                            "buffer_state": None
+                        }
+                    # create tree root node
+                    obs = self.env.set_state(env_state)
+                    tree_node = Node(
+                        state=env_state,
+                        obs=obs,
+                        reward=0,
+                        done=False,
+                        action=None,
+                        parent=RootParentNode(env=self.env),
+                        mcts=self.mcts)
+                else:
+                    # otherwise get last root node from previous time step
+                    tree_node = episode.user_data["tree_node"]
+
+                # run monte carlo simulations to compute the actions
+                # and record the tree
+                mcts_policy, action, tree_node = self.mcts.compute_action(
+                    tree_node)
+                # record action
+                actions.append(action)
+                # store new node
+                episode.user_data["tree_node"] = tree_node
+
+                # store mcts policies vectors and current tree root node
+                if episode.length == 0:
+                    episode.user_data["mcts_policies"] = [mcts_policy]
+                else:
+                    episode.user_data["mcts_policies"].append(mcts_policy)
+
+            return np.array(actions), [], self.extra_action_out(
+                input_dict, state_batches, self.model)
+
+    @override(Policy)
+    def postprocess_trajectory(self,
+                               sample_batch,
+                               other_agent_batches=None,
+                               episode=None):
+        # add mcts policies to sample batch
+        sample_batch["mcts_policies"] = np.array(
+            episode.user_data["mcts_policies"])[sample_batch["t"]]
+        # final episode reward corresponds to the value (if not discounted)
+        # for all transitions in episode
+        final_reward = sample_batch["rewards"][-1]
+        # if r2 is enabled, then add the reward to the buffer and normalize it
+        if self.env.__class__.__name__ == "RankedRewardsEnvWrapper":
+            self.env.r2_buffer.add_reward(final_reward)
+            final_reward = self.env.r2_buffer.normalize(final_reward)
+        sample_batch["value_label"] = final_reward * np.ones_like(
+            sample_batch["t"])
+        return sample_batch
+
+    @override(Policy)
+    def learn_on_batch(self, postprocessed_batch):
+        train_batch = self._lazy_tensor_dict(postprocessed_batch)
+
+        loss_out, policy_loss, value_loss = self._loss(
+            self, self.model, self.dist_class, train_batch)
+        self._optimizer.zero_grad()
+        loss_out.backward()
+
+        grad_process_info = self.extra_grad_process()
+        self._optimizer.step()
+
+        grad_info = self.extra_grad_info(train_batch)
+        grad_info.update(grad_process_info)
+        grad_info.update({
+            "total_loss": loss_out.detach().numpy(),
+            "policy_loss": policy_loss.detach().numpy(),
+            "value_loss": value_loss.detach().numpy()
+        })
+
+        return {LEARNER_STATS_KEY: grad_info}

rllib/contrib/alpha_zero/core/alpha_zero_trainer.py

@@ -0,0 +1,175 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+
+import torch
+import torch.nn as nn
+from ray.rllib.agents import with_common_config
+from ray.rllib.agents.trainer_template import build_trainer
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.models.model import restore_original_dimensions
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.optimizers import SyncSamplesOptimizer
+from ray.rllib.utils import try_import_tf
+from ray.tune.registry import ENV_CREATOR, _global_registry
+
+from ray.rllib.contrib.alpha_zero.core.alpha_zero_policy import AlphaZeroPolicy
+from ray.rllib.contrib.alpha_zero.core.mcts import MCTS
+from ray.rllib.contrib.alpha_zero.core.ranked_rewards import get_r2_env_wrapper
+from ray.rllib.contrib.alpha_zero.optimizer.sync_batches_replay_optimizer \
+    import SyncBatchesReplayOptimizer
+
+tf = try_import_tf()
+logger = logging.getLogger(__name__)
+
+
+def on_episode_start(info):
+    # save env state when an episode starts
+    env = info["env"].get_unwrapped()[0]
+    state = env.get_state()
+    episode = info["episode"]
+    episode.user_data["initial_state"] = state
+
+
+# yapf: disable
+# __sphinx_doc_begin__
+DEFAULT_CONFIG = with_common_config({
+    # Size of batches collected from each worker
+    "sample_batch_size": 200,
+    # Number of timesteps collected for each SGD round
+    "train_batch_size": 4000,
+    # Total SGD batch size across all devices for SGD
+    "sgd_minibatch_size": 128,
+    # Whether to shuffle sequences in the batch when training (recommended)
+    "shuffle_sequences": True,
+    # Number of SGD iterations in each outer loop
+    "num_sgd_iter": 30,
+    # IN case a buffer optimizer is used
+    "learning_starts": 1000,
+    "buffer_size": 10000,
+    # Stepsize of SGD
+    "lr": 5e-5,
+    # Learning rate schedule
+    "lr_schedule": None,
+    # Share layers for value function. If you set this to True, it"s important
+    # to tune vf_loss_coeff.
+    "vf_share_layers": False,
+    # Whether to rollout "complete_episodes" or "truncate_episodes"
+    "batch_mode": "complete_episodes",
+    # Which observation filter to apply to the observation
+    "observation_filter": "NoFilter",
+    # Uses the sync samples optimizer instead of the multi-gpu one. This does
+    # not support minibatches.
+    "simple_optimizer": True,
+
+    # === MCTS ===
+    "mcts_config": {
+        "puct_coefficient": 1.0,
+        "num_simulations": 30,
+        "temperature": 1.5,
+        "dirichlet_epsilon": 0.25,
+        "dirichlet_noise": 0.03,
+        "argmax_tree_policy": False,
+        "add_dirichlet_noise": True,
+    },
+
+    # === Ranked Rewards ===
+    # implement the ranked reward (r2) algorithm
+    # from: https://arxiv.org/pdf/1807.01672.pdf
+    "ranked_rewards": {
+        "enable": True,
+        "percentile": 75,
+        "buffer_max_length": 1000,
+        # add rewards obtained from random policy to
+        # "warm start" the buffer
+        "initialize_buffer": True,
+        "num_init_rewards": 100,
+    },
+
+    # === Evaluation ===
+    # Extra configuration that disables exploration.
+    "evaluation_config": {
+        "mcts_config": {
+            "argmax_tree_policy": True,
+            "add_dirichlet_noise": False,
+        },
+    },
+
+    # === Callbacks ===
+    "callbacks": {
+        "on_episode_start": on_episode_start,
+    }
+})
+
+
+# __sphinx_doc_end__
+# yapf: enable
+
+
+def choose_policy_optimizer(workers, config):
+    if config["simple_optimizer"]:
+        return SyncSamplesOptimizer(
+            workers,
+            num_sgd_iter=config["num_sgd_iter"],
+            train_batch_size=config["train_batch_size"])
+    else:
+        return SyncBatchesReplayOptimizer(
+            workers,
+            num_gradient_descents=config["num_sgd_iter"],
+            learning_starts=config["learning_starts"],
+            train_batch_size=config["train_batch_size"],
+            buffer_size=config["buffer_size"])
+
+
+def alpha_zero_loss(policy, model, dist_class, train_batch):
+    # get inputs unflattened inputs
+    input_dict = restore_original_dimensions(train_batch["obs"],
+                                             policy.observation_space, "torch")
+    # forward pass in model
+    model_out = model.forward(input_dict, None, [1])
+    logits, _ = model_out
+    values = model.value_function()
+    logits, values = torch.squeeze(logits), torch.squeeze(values)
+    priors = nn.Softmax(dim=-1)(logits)
+    # compute actor and critic losses
+    policy_loss = torch.mean(
+        -torch.sum(train_batch["mcts_policies"] * torch.log(priors), dim=-1))
+    value_loss = torch.mean(torch.pow(values - train_batch["value_label"], 2))
+    # compute total loss
+    total_loss = (policy_loss + value_loss) / 2
+    return total_loss, policy_loss, value_loss
+
+
+class AlphaZeroPolicyWrapperClass(AlphaZeroPolicy):
+    def __init__(self, obs_space, action_space, config):
+        model = ModelCatalog.get_model_v2(
+            obs_space, action_space, action_space.n, config["model"], "torch")
+        env_creator = _global_registry.get(ENV_CREATOR, config["env"])
+        if config["ranked_rewards"]["enable"]:
+            # if r2 is enabled, tne env is wrapped to include a rewards buffer
+            # used to normalize rewards
+            env_cls = get_r2_env_wrapper(env_creator, config["ranked_rewards"])
+
+            # the wrapped env is used only in the mcts, not in the
+            # rollout workers
+            def _env_creator():
+                return env_cls(config["env_config"])
+        else:
+
+            def _env_creator():
+                return env_creator(config["env_config"])
+
+        def mcts_creator():
+            return MCTS(model, config["mcts_config"])
+
+        super().__init__(obs_space, action_space, model, alpha_zero_loss,
+                         TorchCategorical, mcts_creator, _env_creator)
+
+
+AlphaZeroTrainer = build_trainer(
+    name="AlphaZero",
+    default_config=DEFAULT_CONFIG,
+    default_policy=AlphaZeroPolicyWrapperClass,
+    make_policy_optimizer=choose_policy_optimizer)

rllib/contrib/alpha_zero/core/mcts.py

@@ -0,0 +1,152 @@
+"""
+Mcts implementation modified from
+https://github.com/brilee/python_uct/blob/master/numpy_impl.py
+"""
+import collections
+import math
+
+import numpy as np
+
+
+class Node:
+    def __init__(self, action, obs, done, reward, state, mcts, parent=None):
+        self.env = parent.env
+        self.action = action  # Action used to go to this state
+
+        self.is_expanded = False
+        self.parent = parent
+        self.children = {}
+
+        self.action_space_size = self.env.action_space.n
+        self.child_total_value = np.zeros(
+            [self.action_space_size], dtype=np.float32)  # Q
+        self.child_priors = np.zeros(
+            [self.action_space_size], dtype=np.float32)  # P
+        self.child_number_visits = np.zeros(
+            [self.action_space_size], dtype=np.float32)  # N
+        self.valid_actions = obs["action_mask"].astype(np.bool)
+
+        self.reward = reward
+        self.done = done
+        self.state = state
+        self.obs = obs
+
+        self.mcts = mcts
+
+    @property
+    def number_visits(self):
+        return self.parent.child_number_visits[self.action]
+
+    @number_visits.setter
+    def number_visits(self, value):
+        self.parent.child_number_visits[self.action] = value
+
+    @property
+    def total_value(self):
+        return self.parent.child_total_value[self.action]
+
+    @total_value.setter
+    def total_value(self, value):
+        self.parent.child_total_value[self.action] = value
+
+    def child_Q(self):
+        # TODO (weak todo) add "softmax" version of the Q-value
+        return self.child_total_value / (1 + self.child_number_visits)
+
+    def child_U(self):
+        return math.sqrt(self.number_visits) * self.child_priors / (
+            1 + self.child_number_visits)
+
+    def best_action(self):
+        """
+        :return: action
+        """
+        child_score = self.child_Q() + self.mcts.c_puct * self.child_U()
+        masked_child_score = child_score
+        masked_child_score[~self.valid_actions] = -np.inf
+        return np.argmax(masked_child_score)
+
+    def select(self):
+        current_node = self
+        while current_node.is_expanded:
+            best_action = current_node.best_action()
+            current_node = current_node.get_child(best_action)
+        return current_node
+
+    def expand(self, child_priors):
+        self.is_expanded = True
+        self.child_priors = child_priors
+
+    def get_child(self, action):
+        if action not in self.children:
+            self.env.set_state(self.state)
+            obs, reward, done, _ = self.env.step(action)
+            next_state = self.env.get_state()
+            self.children[action] = Node(
+                state=next_state,
+                action=action,
+                parent=self,
+                reward=reward,
+                done=done,
+                obs=obs,
+                mcts=self.mcts)
+        return self.children[action]
+
+    def backup(self, value):
+        current = self
+        while current.parent is not None:
+            current.number_visits += 1
+            current.total_value += value
+            current = current.parent
+
+
+class RootParentNode(object):
+    def __init__(self, env):
+        self.parent = None
+        self.child_total_value = collections.defaultdict(float)
+        self.child_number_visits = collections.defaultdict(float)
+        self.env = env
+
+
+class MCTS:
+    def __init__(self, model, mcts_param):
+        self.model = model
+        self.temperature = mcts_param["temperature"]
+        self.dir_epsilon = mcts_param["dirichlet_epsilon"]
+        self.dir_noise = mcts_param["dirichlet_noise"]
+        self.num_sims = mcts_param["num_simulations"]
+        self.exploit = mcts_param["argmax_tree_policy"]
+        self.add_dirichlet_noise = mcts_param["add_dirichlet_noise"]
+        self.c_puct = mcts_param["puct_coefficient"]
+
+    def compute_action(self, node):
+        for _ in range(self.num_sims):
+            leaf = node.select()
+            if leaf.done:
+                value = leaf.reward
+            else:
+                child_priors, value = self.model.compute_priors_and_value(
+                    leaf.obs)
+                if self.add_dirichlet_noise:
+                    child_priors = (1 - self.dir_epsilon) * child_priors
+                    child_priors += self.dir_epsilon * np.random.dirichlet(
+                        [self.dir_noise] * child_priors.size)
+
+                leaf.expand(child_priors)
+            leaf.backup(value)
+
+        # Tree policy target (TPT)
+        tree_policy = node.child_number_visits / node.number_visits
+        tree_policy = tree_policy / np.max(
+            tree_policy)  # to avoid overflows when computing softmax
+        tree_policy = np.power(tree_policy, self.temperature)
+        tree_policy = tree_policy / np.sum(tree_policy)
+        if self.exploit:
+            # if exploit then choose action that has the maximum
+            # tree policy probability
+            action = np.argmax(tree_policy)
+        else:
+            # otherwise sample an action according to tree policy probabilities
+            action = np.random.choice(
+                np.arange(node.action_space_size), p=tree_policy)
+        return tree_policy, action, node.children[action]

rllib/contrib/alpha_zero/core/ranked_rewards.py

@@ -0,0 +1,79 @@
+from copy import deepcopy
+
+import numpy as np
+
+
+class RankedRewardsBuffer:
+    def __init__(self, buffer_max_length, percentile):
+        self.buffer_max_length = buffer_max_length
+        self.percentile = percentile
+        self.buffer = []
+
+    def add_reward(self, reward):
+        if len(self.buffer) < self.buffer_max_length:
+            self.buffer.append(reward)
+        else:
+            self.buffer = self.buffer[1:] + [reward]
+
+    def normalize(self, reward):
+        reward_threshold = np.percentile(self.buffer, self.percentile)
+        if reward < reward_threshold:
+            return -1.0
+        else:
+            return 1.0
+
+    def get_state(self):
+        return np.array(self.buffer)
+
+    def set_state(self, state):
+        if state is not None:
+            self.buffer = list(state)
+
+
+def get_r2_env_wrapper(env_creator, r2_config):
+    class RankedRewardsEnvWrapper:
+        def __init__(self, env_config):
+            self.env = env_creator(env_config)
+            self.action_space = self.env.action_space
+            self.observation_space = self.env.observation_space
+            max_buffer_length = r2_config["buffer_max_length"]
+            percentile = r2_config["percentile"]
+            self.r2_buffer = RankedRewardsBuffer(max_buffer_length, percentile)
+            if r2_config["initialize_buffer"]:
+                self._initialize_buffer(r2_config["num_init_rewards"])
+
+        def _initialize_buffer(self, num_init_rewards=100):
+            # initialize buffer with random policy
+            for _ in range(num_init_rewards):
+                obs = self.env.reset()
+                done = False
+                while not done:
+                    mask = obs["action_mask"]
+                    probs = mask / mask.sum()
+                    action = np.random.choice(
+                        np.arange(mask.shape[0]), p=probs)
+                    obs, reward, done, _ = self.env.step(action)
+                self.r2_buffer.add_reward(reward)
+
+        def step(self, action):
+            obs, reward, done, info = self.env.step(action)
+            if done:
+                reward = self.r2_buffer.normalize(reward)
+            return obs, reward, done, info
+
+        def get_state(self):
+            state = {
+                "env_state": self.env.get_state(),
+                "buffer_state": self.r2_buffer.get_state()
+            }
+            return deepcopy(state)
+
+        def reset(self):
+            return self.env.reset()
+
+        def set_state(self, state):
+            obs = self.env.set_state(state["env_state"])
+            self.r2_buffer.set_state(state["buffer_state"])
+            return obs
+
+    return RankedRewardsEnvWrapper

rllib/contrib/alpha_zero/environments/cartpole.py

@@ -0,0 +1,40 @@
+from copy import deepcopy
+
+import gym
+import numpy as np
+from gym.spaces import Discrete, Dict, Box
+
+
+class CartPole:
+    """
+    Wrapper for gym CartPole environment where the reward
+    is accumulated to the end
+    """
+
+    def __init__(self, config=None):
+        self.env = gym.make("CartPole-v0")
+        self.action_space = Discrete(2)
+        self.observation_space = Dict({
+            "obs": self.env.observation_space,
+            "action_mask": Box(low=0, high=1, shape=(self.action_space.n, ))
+        })
+        self.running_reward = 0
+
+    def reset(self):
+        self.running_reward = 0
+        return {"obs": self.env.reset(), "action_mask": np.array([1, 1])}
+
+    def step(self, action):
+        obs, rew, done, info = self.env.step(action)
+        self.running_reward += rew
+        score = self.running_reward if done else 0
+        return {"obs": obs, "action_mask": np.array([1, 1])}, score, done, info
+
+    def set_state(self, state):
+        self.running_reward = state[1]
+        self.env = deepcopy(state[0])
+        obs = np.array(list(self.env.unwrapped.state))
+        return {"obs": obs, "action_mask": np.array([1, 1])}
+
+    def get_state(self):
+        return deepcopy(self.env), self.running_reward

rllib/contrib/alpha_zero/examples/train_cartpole.py

@@ -0,0 +1,51 @@
+"""Example of using training on CartPole."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+
+from ray import tune
+
+from ray.rllib.contrib.alpha_zero.models.custom_torch_models import DenseModel
+from ray.rllib.contrib.alpha_zero.environments.cartpole import CartPole
+from ray.rllib.models.catalog import ModelCatalog
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-workers", default=6, type=int)
+    parser.add_argument("--training-iteration", default=10000, type=int)
+    args = parser.parse_args()
+
+    ModelCatalog.register_custom_model("dense_model", DenseModel)
+
+    tune.run(
+        "contrib/AlphaZero",
+        stop={"training_iteration": args.training_iteration},
+        max_failures=0,
+        config={
+            "env": CartPole,
+            "num_workers": args.num_workers,
+            "sample_batch_size": 50,
+            "train_batch_size": 500,
+            "sgd_minibatch_size": 64,
+            "lr": 1e-4,
+            "num_sgd_iter": 1,
+            "mcts_config": {
+                "puct_coefficient": 1.5,
+                "num_simulations": 100,
+                "temperature": 1.0,
+                "dirichlet_epsilon": 0.20,
+                "dirichlet_noise": 0.03,
+                "argmax_tree_policy": False,
+                "add_dirichlet_noise": True,
+            },
+            "ranked_rewards": {
+                "enable": True,
+            },
+            "model": {
+                "custom_model": "dense_model",
+            },
+        },
+    )

rllib/contrib/alpha_zero/models/custom_torch_models.py

@@ -0,0 +1,108 @@
+from abc import ABC
+
+import numpy as np
+import torch
+import torch.nn as nn
+from ray.rllib.models.model import restore_original_dimensions
+from ray.rllib.models.preprocessors import get_preprocessor
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+
+
+def convert_to_tensor(arr):
+    tensor = torch.from_numpy(np.asarray(arr))
+    if tensor.dtype == torch.double:
+        tensor = tensor.float()
+    return tensor
+
+
+class ActorCriticModel(TorchModelV2, nn.Module, ABC):
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
+                              model_config, name)
+        nn.Module.__init__(self)
+
+        self.preprocessor = get_preprocessor(obs_space.original_space)(
+            obs_space.original_space)
+
+        self.shared_layers = None
+        self.actor_layers = None
+        self.critic_layers = None
+
+        self._value_out = None
+
+    def forward(self, input_dict, state, seq_lens):
+        x = input_dict["obs"]
+        x = self.shared_layers(x)
+        # actor outputs
+        logits = self.actor_layers(x)
+
+        # compute value
+        self._value_out = self.critic_layers(x)
+        return logits, None
+
+    def value_function(self):
+        return self._value_out
+
+    def compute_priors_and_value(self, obs):
+        obs = convert_to_tensor([self.preprocessor.transform(obs)])
+        input_dict = restore_original_dimensions(obs, self.obs_space, "torch")
+
+        with torch.no_grad():
+            model_out = self.forward(input_dict, None, [1])
+            logits, _ = model_out
+            value = self.value_function()
+            logits, value = torch.squeeze(logits), torch.squeeze(value)
+            priors = nn.Softmax(dim=-1)(logits)
+
+            priors = priors.cpu().numpy()
+            value = value.cpu().numpy()
+
+            return priors, value
+
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+class ConvNetModel(ActorCriticModel):
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        ActorCriticModel.__init__(self, obs_space, action_space, num_outputs,
+                                  model_config, name)
+
+        in_channels = model_config["custom_options"]["in_channels"]
+        feature_dim = model_config["custom_options"]["feature_dim"]
+
+        self.shared_layers = nn.Sequential(
+            nn.Conv2d(in_channels, 32, kernel_size=4, stride=2),
+            nn.Conv2d(32, 64, kernel_size=2, stride=1),
+            nn.Conv2d(64, 64, kernel_size=2, stride=1), Flatten(),
+            nn.Linear(1024, feature_dim))
+
+        self.actor_layers = nn.Sequential(
+            nn.Linear(in_features=feature_dim, out_features=action_space.n))
+
+        self.critic_layers = nn.Sequential(
+            nn.Linear(in_features=feature_dim, out_features=1))
+
+        self._value_out = None
+
+
+class DenseModel(ActorCriticModel):
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        ActorCriticModel.__init__(self, obs_space, action_space, num_outputs,
+                                  model_config, name)
+
+        self.shared_layers = nn.Sequential(
+            nn.Linear(
+                in_features=obs_space.original_space["obs"].shape[0],
+                out_features=256), nn.Linear(
+                    in_features=256, out_features=256))
+        self.actor_layers = nn.Sequential(
+            nn.Linear(in_features=256, out_features=action_space.n))
+        self.critic_layers = nn.Sequential(
+            nn.Linear(in_features=256, out_features=1))
+        self._value_out = None

rllib/contrib/alpha_zero/optimizer/sync_batches_replay_optimizer.py

@@ -0,0 +1,34 @@
+import random
+
+from ray.rllib.evaluation.metrics import get_learner_stats
+from ray.rllib.optimizers.sync_batch_replay_optimizer import \
+    SyncBatchReplayOptimizer
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+
+
+class SyncBatchesReplayOptimizer(SyncBatchReplayOptimizer):
+    def __init__(self,
+                 workers,
+                 learning_starts=1000,
+                 buffer_size=10000,
+                 train_batch_size=32,
+                 num_gradient_descents=10):
+        super(SyncBatchesReplayOptimizer, self).__init__(
+            workers, learning_starts, buffer_size, train_batch_size)
+        self.num_sgds = num_gradient_descents
+
+    @override(SyncBatchReplayOptimizer)
+    def _optimize(self):
+        for _ in range(self.num_sgds):
+            samples = [random.choice(self.replay_buffer)]
+            while sum(s.count for s in samples) < self.train_batch_size:
+                samples.append(random.choice(self.replay_buffer))
+            samples = SampleBatch.concat_samples(samples)
+            with self.grad_timer:
+                info_dict = self.workers.local_worker().learn_on_batch(samples)
+                for policy_id, info in info_dict.items():
+                    self.learner_stats[policy_id] = get_learner_stats(info)
+                self.grad_timer.push_units_processed(samples.count)
+            self.num_steps_trained += samples.count
+        return info_dict

rllib/contrib/registry.py

@@ -15,7 +15,14 @@ def _import_maddpg():
     return maddpg.MADDPGTrainer
 
 
+def _import_alphazero():
+    from ray.rllib.contrib.alpha_zero.core.alpha_zero_trainer import\
+        AlphaZeroTrainer
+    return AlphaZeroTrainer
+
+
 CONTRIBUTED_ALGORITHMS = {
     "contrib/RandomAgent": _import_random_agent,
     "contrib/MADDPG": _import_maddpg,
+    "contrib/AlphaZero": _import_alphazero,
 }