From 6e6c680f145e1595ffa55364f77e349a96fadfe3 Mon Sep 17 00:00:00 2001 From: Michael Luo Date: Thu, 12 Nov 2020 10:30:41 -0800 Subject: [PATCH] MBMPO Cartpole (#11832) * MBMPO Cartpole Done * Added doc --- doc/source/rllib-algorithms.rst | 7 ++++- rllib/agents/mbmpo/model_ensemble.py | 5 ++++ rllib/agents/mbmpo/tests/test_mbmpo.py | 2 +- rllib/env/model_vector_env.py | 9 ++++++- rllib/examples/env/mbmpo_env.py | 22 ++++++++++++++- .../tuned_examples/mbmpo/cartpole-mbmpo.yaml | 27 +++++++++++++++++++ 6 files changed, 68 insertions(+), 4 deletions(-) create mode 100644 rllib/tuned_examples/mbmpo/cartpole-mbmpo.yaml diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index 4068dce49..aac235e94 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -471,7 +471,12 @@ RLlib's MBMPO implementation is a Dyna-styled model-based RL method that learns Additional statistics are logged in MBMPO. Each MBMPO iteration corresponds to multiple MAML iterations, and ``MAMLIter$i$_DynaTrajInner_$j$_episode_reward_mean`` measures the agent's returns across the dynamics models at iteration ``i`` of MAML and step ``j`` of inner adaptation. Examples can be seen `here `__. -Tuned examples: `HalfCheetah `__, `Hopper `__ +Tuned examples (continuous actions): +`Pendulum-v0 `__, +`HalfCheetah `__, +`Hopper `__, +Tuned examples (discrete actions): +`CartPole-v0 `__ **MuJoCo results @100K steps:** `more details `__ diff --git a/rllib/agents/mbmpo/model_ensemble.py b/rllib/agents/mbmpo/model_ensemble.py index 72749eedc..2bb9513da 100644 --- a/rllib/agents/mbmpo/model_ensemble.py +++ b/rllib/agents/mbmpo/model_ensemble.py @@ -206,6 +206,11 @@ class DynamicsEnsembleCustomModel(TorchModelV2, nn.Module): # Process Samples new_samples = process_samples(new_samples) + if isinstance(self.action_space, Discrete): + act = new_samples["actions"] + new_act = np.zeros((act.size, act.max() + 1)) + new_act[np.arange(act.size), act] = 1 + new_samples["actions"] = new_act.astype("float32") if not self.replay_buffer: self.replay_buffer = new_samples diff --git a/rllib/agents/mbmpo/tests/test_mbmpo.py b/rllib/agents/mbmpo/tests/test_mbmpo.py index 3deecd5c7..de708fd50 100644 --- a/rllib/agents/mbmpo/tests/test_mbmpo.py +++ b/rllib/agents/mbmpo/tests/test_mbmpo.py @@ -27,7 +27,7 @@ class TestMBMPO(unittest.TestCase): for _ in framework_iterator(config, frameworks="torch"): trainer = mbmpo.MBMPOTrainer( config=config, - env="ray.rllib.examples.env.mbmpo_env.PendulumWrapper") + env="ray.rllib.examples.env.mbmpo_env.CartPoleWrapper") for i in range(num_iterations): trainer.train() check_compute_single_action( diff --git a/rllib/env/model_vector_env.py b/rllib/env/model_vector_env.py index 0b1b7763e..b5b9700d6 100644 --- a/rllib/env/model_vector_env.py +++ b/rllib/env/model_vector_env.py @@ -1,6 +1,6 @@ import logging import numpy as np - +from gym.spaces import Discrete from ray.rllib.utils.annotations import override from ray.rllib.env.vector_env import VectorEnv from ray.rllib.evaluation.rollout_worker import get_global_worker @@ -94,6 +94,13 @@ class _VectorizedModelGymEnv(VectorEnv): if self.cur_obs is None: raise ValueError("Need to reset env first") + # If discrete, need to one-hot actions + if isinstance(self.action_space, Discrete): + act = np.array(actions) + new_act = np.zeros((act.size, act.max() + 1)) + new_act[np.arange(act.size), act] = 1 + actions = new_act.astype("float32") + # Batch the TD-model prediction. obs_batch = np.stack(self.cur_obs, axis=0) action_batch = np.stack(actions, axis=0) diff --git a/rllib/examples/env/mbmpo_env.py b/rllib/examples/env/mbmpo_env.py index 947a34fd3..c49ef77be 100644 --- a/rllib/examples/env/mbmpo_env.py +++ b/rllib/examples/env/mbmpo_env.py @@ -1,5 +1,5 @@ import gym -from gym.envs.classic_control import PendulumEnv +from gym.envs.classic_control import PendulumEnv, CartPoleEnv import numpy as np # MuJoCo may not be installed. @@ -10,6 +10,26 @@ except (ImportError, gym.error.DependencyNotInstalled): pass +class CartPoleWrapper(CartPoleEnv): + """Wrapper for the Cartpole-v0 environment. + + Adds an additional `reward` method for some model-based RL algos (e.g. + MB-MPO). + """ + + def reward(self, obs, action, obs_next): + # obs = batch * [pos, vel, angle, rotation_rate] + x = obs_next[:, 0] + theta = obs_next[:, 2] + + rew = (x < -self.x_threshold) | (x > self.x_threshold) | ( + theta < -self.theta_threshold_radians) | ( + theta > self.theta_threshold_radians) + + rew = rew.astype(float) + return rew + + class PendulumWrapper(PendulumEnv): """Wrapper for the Pendulum-v0 environment. diff --git a/rllib/tuned_examples/mbmpo/cartpole-mbmpo.yaml b/rllib/tuned_examples/mbmpo/cartpole-mbmpo.yaml new file mode 100644 index 000000000..604efd8b3 --- /dev/null +++ b/rllib/tuned_examples/mbmpo/cartpole-mbmpo.yaml @@ -0,0 +1,27 @@ +cartpole-mbmpo: + env: ray.rllib.examples.env.mbmpo_env.CartPoleWrapper + run: MBMPO + stop: + episode_reward_mean: 190 + training_iteration: 20 + config: + # Only supported in torch right now. + framework: torch + #horizon: 200 + num_envs_per_worker: 20 + inner_adaptation_steps: 1 + maml_optimizer_steps: 8 + gamma: 0.99 + lambda: 1.0 + lr: 0.001 + clip_param: 0.5 + kl_target: 0.003 + kl_coeff: 0.0000000001 + num_workers: 10 + num_gpus: 0 + inner_lr: 0.001 + clip_actions: False + num_maml_steps: 15 + model: + fcnet_hiddens: [32, 32] + free_log_std: True