Chapter 12: Meta RL

Author: ebilgin

Chapter12/ml-test.py

@@ -0,0 +1,50 @@
+from copy import deepcopy
+
+import ray
+from ray import tune
+
+#from ray.rllib.examples.env.pendulum_mass import PendulumMassEnv
+from ray.rllib.agents.maml.maml import MAMLTrainer, DEFAULT_CONFIG
+
+import numpy as np
+import gym
+from gym.envs.classic_control.pendulum import PendulumEnv
+from ray.rllib.env.meta_env import MetaEnv
+
+from penenv3 import PenEnv
+
+config = deepcopy(DEFAULT_CONFIG)
+
+
+ray.init()
+tune.run(
+    "MAML",
+    stop={"training_iteration": 500},
+    config=dict(
+        DEFAULT_CONFIG,
+        **{
+            "env": PenEnv,
+            "horizon": 200,
+            "rollout_fragment_length": 200,
+            "num_envs_per_worker": 10,
+            "inner_adaptation_steps": 1,
+            "maml_optimizer_steps": 5,
+            "gamma": 0.99,
+            "lambda": 1.0,
+            "lr": 0.001,
+            "vf_loss_coeff": 0.5,
+            "clip_param": 0.3,
+            "kl_target": 0.01,
+            "kl_coeff": 0.001,
+            "num_workers": 60,
+            "num_gpus": 1,
+            "inner_lr": 0.03,
+            "explore": True,
+            "clip_actions": False,
+            "model": {"fcnet_hiddens": [64, 64],
+                      "free_log_std": True
+                      }
+        }
+    ),
+    checkpoint_freq=10,
+)

Chapter12/ml-test_lstm.py

@@ -0,0 +1,49 @@
+from copy import deepcopy
+
+import ray
+from ray import tune
+
+from ray.rllib.agents.ppo.ppo import PPOTrainer, DEFAULT_CONFIG
+
+import numpy as np
+import gym
+from gym.envs.classic_control.pendulum import PendulumEnv
+from ray.rllib.env.meta_env import MetaEnv
+
+from penenv2 import PenEnv2
+
+config = deepcopy(DEFAULT_CONFIG)
+
+
+ray.init()
+tune.run(
+    PPOTrainer,
+    #stop={"training_iteration": 500},
+    config=dict(
+        DEFAULT_CONFIG,
+        **{
+            "env": PenEnv2,
+            "horizon": 200,
+            "rollout_fragment_length": 200,
+            #"num_envs_per_worker": 10,
+            "gamma": 0.99,
+            "lambda": 1.0,
+            "lr": 0.001,
+            "vf_loss_coeff": 0.5,
+            "clip_param": 0.3,
+            "kl_target": 0.01,
+            "kl_coeff": 0.001,
+            "num_workers": 60,
+            "num_gpus": 1,
+
+            "clip_actions": False,
+            "model": {#"fcnet_hiddens": [64, 64],
+                        "use_lstm": True,
+                        "lstm_cell_size": 128,
+                        "lstm_use_prev_action_reward": True,
+                        "max_seq_len": 10
+                      }
+        }
+    ),
+    checkpoint_freq=10,
+)

Chapter12/penenv2.py

@@ -0,0 +1,128 @@
+import gym
+from gym import spaces
+from gym.utils import seeding
+import numpy as np
+from os import path
+
+
+import gym
+from gym import spaces
+from gym.utils import seeding
+import numpy as np
+from os import path
+
+
+class PenEnv2(gym.Env):
+    metadata = {
+        'render.modes': ['human', 'rgb_array'],
+        'video.frames_per_second': 30
+    }
+
+    def __init__(self, env_config={}):
+        self.max_speed = 8
+        self.max_torque = 2.
+        self.dt = .05
+        self.g = env_config.get("g", 10)
+        self.m = 1.
+        self.l = 1.
+        self.viewer = None
+
+        high = np.array([1., 1., self.max_speed], dtype=np.float32)
+        self.action_space = spaces.Box(
+            low=-self.max_torque,
+            high=self.max_torque, shape=(1,),
+            dtype=np.float32
+        )
+        self.observation_space = spaces.Box(
+            low=-high,
+            high=high,
+            dtype=np.float32
+        )
+
+        self.seed()
+
+    def seed(self, seed=None):
+        self.np_random, seed = seeding.np_random(seed)
+        return [seed]
+
+    def step(self, u):
+        th, thdot = self.state  # th := theta
+
+        g = self.g
+        m = self.m
+        l = self.l
+        dt = self.dt
+
+        u = np.clip(u, -self.max_torque, self.max_torque)[0]
+        self.last_u = u  # for rendering
+        costs = angle_normalize(th) ** 2 + .1 * thdot ** 2 + .001 * (u ** 2)
+
+        newthdot = thdot + (-3 * g / (2 * l) * np.sin(th + np.pi) + 3. / (m * l ** 2) * u) * dt
+        newth = th + newthdot * dt
+        newthdot = np.clip(newthdot, -self.max_speed, self.max_speed)
+
+        self.state = np.array([newth, newthdot])
+        return self._get_obs(), -costs, False, {}
+
+    def reset(self):
+        high = np.array([np.pi, 1])
+        self.state = self.np_random.uniform(low=-high, high=high)
+        self.last_u = None
+        self.m = np.random.uniform(low=0.5, high=2.0)
+        return self._get_obs()
+
+    def _get_obs(self):
+        theta, thetadot = self.state
+        return np.array([np.cos(theta), np.sin(theta), thetadot])
+
+    def render(self, mode='human'):
+        if self.viewer is None:
+            from gym.envs.classic_control import rendering
+            self.viewer = rendering.Viewer(500, 500)
+            self.viewer.set_bounds(-2.2, 2.2, -2.2, 2.2)
+            rod = rendering.make_capsule(1, .2)
+            rod.set_color(.8, .3, .3)
+            self.pole_transform = rendering.Transform()
+            rod.add_attr(self.pole_transform)
+            self.viewer.add_geom(rod)
+            axle = rendering.make_circle(.05)
+            axle.set_color(0, 0, 0)
+            self.viewer.add_geom(axle)
+            fname = path.join(path.dirname(__file__), "assets/clockwise.png")
+            self.img = rendering.Image(fname, 1., 1.)
+            self.imgtrans = rendering.Transform()
+            self.img.add_attr(self.imgtrans)
+
+        self.viewer.add_onetime(self.img)
+        self.pole_transform.set_rotation(self.state[0] + np.pi / 2)
+        if self.last_u:
+            self.imgtrans.scale = (-self.last_u / 2, np.abs(self.last_u) / 2)
+
+        return self.viewer.render(return_rgb_array=mode == 'rgb_array')
+
+    def close(self):
+        if self.viewer:
+            self.viewer.close()
+            self.viewer = None
+
+    def sample_tasks(self, n_tasks):
+        # Mass is a random float between 0.5 and 2
+        return np.random.uniform(low=0.5, high=2.0, size=(n_tasks, ))
+
+    def set_task(self, task):
+        """
+        Args:
+            task: task of the meta-learning environment
+        """
+        self.m = task
+
+    def get_task(self):
+        """
+        Returns:
+            task: task of the meta-learning environment
+        """
+        return self.m
+
+
+def angle_normalize(x):
+    return (((x+np.pi) % (2*np.pi)) - np.pi)

Chapter12/penenv3.py

@@ -0,0 +1,127 @@
+import gym
+from gym import spaces
+from gym.utils import seeding
+import numpy as np
+from os import path
+
+
+import gym
+from gym import spaces
+from gym.utils import seeding
+import numpy as np
+from os import path
+
+
+class PenEnv(gym.Env):
+    metadata = {
+        'render.modes': ['human', 'rgb_array'],
+        'video.frames_per_second': 30
+    }
+
+    def __init__(self, env_config={}):
+        self.max_speed = 8
+        self.max_torque = 2.
+        self.dt = .05
+        self.g = env_config.get("g", 10)
+        self.m = 1.
+        self.l = 1.
+        self.viewer = None
+
+        high = np.array([1., 1., self.max_speed], dtype=np.float32)
+        self.action_space = spaces.Box(
+            low=-self.max_torque,
+            high=self.max_torque, shape=(1,),
+            dtype=np.float32
+        )
+        self.observation_space = spaces.Box(
+            low=-high,
+            high=high,
+            dtype=np.float32
+        )
+
+        self.seed()
+
+    def seed(self, seed=None):
+        self.np_random, seed = seeding.np_random(seed)
+        return [seed]
+
+    def step(self, u):
+        th, thdot = self.state  # th := theta
+
+        g = self.g
+        m = self.m
+        l = self.l
+        dt = self.dt
+
+        u = np.clip(u, -self.max_torque, self.max_torque)[0]
+        self.last_u = u  # for rendering
+        costs = angle_normalize(th) ** 2 + .1 * thdot ** 2 + .001 * (u ** 2)
+
+        newthdot = thdot + (-3 * g / (2 * l) * np.sin(th + np.pi) + 3. / (m * l ** 2) * u) * dt
+        newth = th + newthdot * dt
+        newthdot = np.clip(newthdot, -self.max_speed, self.max_speed)
+
+        self.state = np.array([newth, newthdot])
+        return self._get_obs(), -costs, False, {}
+
+    def reset(self):
+        high = np.array([np.pi, 1])
+        self.state = self.np_random.uniform(low=-high, high=high)
+        self.last_u = None
+        return self._get_obs()
+
+    def _get_obs(self):
+        theta, thetadot = self.state
+        return np.array([np.cos(theta), np.sin(theta), thetadot])
+
+    def render(self, mode='human'):
+        if self.viewer is None:
+            from gym.envs.classic_control import rendering
+            self.viewer = rendering.Viewer(500, 500)
+            self.viewer.set_bounds(-2.2, 2.2, -2.2, 2.2)
+            rod = rendering.make_capsule(1, .2)
+            rod.set_color(.8, .3, .3)
+            self.pole_transform = rendering.Transform()
+            rod.add_attr(self.pole_transform)
+            self.viewer.add_geom(rod)
+            axle = rendering.make_circle(.05)
+            axle.set_color(0, 0, 0)
+            self.viewer.add_geom(axle)
+            fname = path.join(path.dirname(__file__), "assets/clockwise.png")
+            self.img = rendering.Image(fname, 1., 1.)
+            self.imgtrans = rendering.Transform()
+            self.img.add_attr(self.imgtrans)
+
+        self.viewer.add_onetime(self.img)
+        self.pole_transform.set_rotation(self.state[0] + np.pi / 2)
+        if self.last_u:
+            self.imgtrans.scale = (-self.last_u / 2, np.abs(self.last_u) / 2)
+
+        return self.viewer.render(return_rgb_array=mode == 'rgb_array')
+
+    def close(self):
+        if self.viewer:
+            self.viewer.close()
+            self.viewer = None
+
+    def sample_tasks(self, n_tasks):
+        # Mass is a random float between 0.5 and 2
+        return np.random.uniform(low=0.5, high=2.0, size=(n_tasks, ))
+
+    def set_task(self, task):
+        """
+        Args:
+            task: task of the meta-learning environment
+        """
+        self.m = task
+
+    def get_task(self):
+        """
+        Returns:
+            task: task of the meta-learning environment
+        """
+        return self.m
+
+
+def angle_normalize(x):
+    return (((x+np.pi) % (2*np.pi)) - np.pi)