diff --git a/ci/jenkins_tests/run_rllib_tests.sh b/ci/jenkins_tests/run_rllib_tests.sh index 07f34b11e..05db3d3ab 100755 --- a/ci/jenkins_tests/run_rllib_tests.sh +++ b/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 diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index 90826674e..8ce7e66ff 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -364,3 +364,17 @@ Tuned examples: `CartPole-v0 `__ `[implementation] `__ 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) `__ 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 `__ + +**AlphaZero-specific configs** (see also `common configs `__): + +.. literalinclude:: ../../rllib/contrib/alpha_zero/core/alpha_zero_trainer.py + :language: python + :start-after: __sphinx_doc_begin__ + :end-before: __sphinx_doc_end__ diff --git a/doc/source/rllib-toc.rst b/doc/source/rllib-toc.rst index fca6a7623..86f74dd3c 100644 --- a/doc/source/rllib-toc.rst +++ b/doc/source/rllib-toc.rst @@ -78,6 +78,8 @@ Algorithms - `Asynchronous Proximal Policy Optimization (APPO) `__ + - `Single-Player AlphaZero (contrib/AlphaZero) `__ + * Gradient-based - `Advantage Actor-Critic (A2C, A3C) `__ diff --git a/rllib/contrib/alpha_zero/README.md b/rllib/contrib/alpha_zero/README.md new file mode 100644 index 000000000..2b9807e29 --- /dev/null +++ b/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 + \ No newline at end of file diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py new file mode 100644 index 000000000..2b1a29059 --- /dev/null +++ b/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} diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py b/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py new file mode 100644 index 000000000..0b3d9a9e9 --- /dev/null +++ b/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) diff --git a/rllib/contrib/alpha_zero/core/mcts.py b/rllib/contrib/alpha_zero/core/mcts.py new file mode 100644 index 000000000..c31538236 --- /dev/null +++ b/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] diff --git a/rllib/contrib/alpha_zero/core/ranked_rewards.py b/rllib/contrib/alpha_zero/core/ranked_rewards.py new file mode 100644 index 000000000..67ff3853a --- /dev/null +++ b/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 diff --git a/rllib/contrib/alpha_zero/doc/cartpole_plot.png b/rllib/contrib/alpha_zero/doc/cartpole_plot.png new file mode 100644 index 000000000..75cbf2146 Binary files /dev/null and b/rllib/contrib/alpha_zero/doc/cartpole_plot.png differ diff --git a/rllib/contrib/alpha_zero/environments/cartpole.py b/rllib/contrib/alpha_zero/environments/cartpole.py new file mode 100644 index 000000000..961d33689 --- /dev/null +++ b/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 diff --git a/rllib/contrib/alpha_zero/examples/train_cartpole.py b/rllib/contrib/alpha_zero/examples/train_cartpole.py new file mode 100644 index 000000000..8b9cbeecd --- /dev/null +++ b/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", + }, + }, + ) diff --git a/rllib/contrib/alpha_zero/models/custom_torch_models.py b/rllib/contrib/alpha_zero/models/custom_torch_models.py new file mode 100644 index 000000000..bf0291ee4 --- /dev/null +++ b/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 diff --git a/rllib/contrib/alpha_zero/optimizer/sync_batches_replay_optimizer.py b/rllib/contrib/alpha_zero/optimizer/sync_batches_replay_optimizer.py new file mode 100644 index 000000000..839808c5b --- /dev/null +++ b/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 diff --git a/rllib/contrib/registry.py b/rllib/contrib/registry.py index 340cd701c..05c456d23 100644 --- a/rllib/contrib/registry.py +++ b/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, }