MBMPO Cartpole (#11832)

* MBMPO Cartpole Done

* Added doc

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 <https://github.com/ray-project/rl-experiments/tree/master/mbmpo>`__.
 
-Tuned examples: `HalfCheetah <https://github.com/ray-project/ray/blob/master/rllib/tuned_examples/mbmpo/halfcheetah-mbmpo.yaml>`__, `Hopper <https://github.com/ray-project/ray/blob/master/rllib/tuned_examples/mbmpo/hopper-mbmpo.yaml>`__
+Tuned examples (continuous actions):
+`Pendulum-v0 <https://github.com/ray-project/ray/blob/master/rllib/tuned_examples/mbmpo/pendulum-mbmpo.yaml>`__,
+`HalfCheetah <https://github.com/ray-project/ray/blob/master/rllib/tuned_examples/mbmpo/halfcheetah-mbmpo.yaml>`__,
+`Hopper <https://github.com/ray-project/ray/blob/master/rllib/tuned_examples/mbmpo/hopper-mbmpo.yaml>`__,
+Tuned examples (discrete actions):
+`CartPole-v0 <https://github.com/ray-project/ray/blob/master/rllib/tuned_examples/mbmpo/cartpole-mbmpo.yaml>`__
 
 **MuJoCo results @100K steps:** `more details <https://github.com/ray-project/rl-experiments>`__
 

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

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(

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)

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.
 

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