Merge pull request #1 from fairinternal/add_hydra

Added hydra configuration for MBPO

Author: luisenp

.gitignore

@@ -1,2 +1,3 @@
 *personal*
-.idea
+.idea
+notebooks/.ipynb_checkpoints

conf/agent/sac.yaml

@@ -0,0 +1,38 @@
+# @package _global_
+agent:
+  _target_: pytorch_sac.agent.sac.SACAgent
+  obs_dim: ??? # to be specified later
+  action_dim: ??? # to be specified later
+  action_range: ??? # to be specified later
+  device: ${device}
+  critic_cfg: ${double_q_critic}
+  actor_cfg: ${diag_gaussian_actor}
+  discount: 0.99
+  init_temperature: 0.1
+  alpha_lr: 1e-4
+  alpha_betas: [0.9, 0.999]
+  actor_lr: 1e-4
+  actor_betas: [0.9, 0.999]
+  actor_update_frequency: 1
+  critic_lr: 1e-4
+  critic_betas: [0.9, 0.999]
+  critic_tau: 0.005
+  critic_target_update_frequency: 2
+  batch_size: 1024
+  learnable_temperature: true
+  target_entropy: -1
+
+double_q_critic:
+  _target_: pytorch_sac.agent.critic.DoubleQCritic
+  obs_dim: ${agent.obs_dim}
+  action_dim: ${agent.action_dim}
+  hidden_dim: 1024
+  hidden_depth: 2
+
+diag_gaussian_actor:
+  _target_: pytorch_sac.agent.actor.DiagGaussianActor
+  obs_dim: ${agent.obs_dim}
+  action_dim: ${agent.action_dim}
+  hidden_depth: 2
+  hidden_dim: 1024
+  log_std_bounds: [-5, 2]

conf/mbpo.yaml

@@ -0,0 +1,49 @@
+defaults:
+  - agent: sac
+
+model:
+  _target_: mbrl.models.Ensemble
+  ensemble_size: 7
+  in_size: ???
+  out_size: ???
+  member_cfg: ${member_cfg}
+  device: ${device}
+  optim_lr: 0.0075 # TODO this should be moved out of the ensemble class
+
+member_cfg:
+  _target_: mbrl.models.GaussianMLP
+  device: ${model.device}
+  num_layers: 4
+  in_size: ${model.in_size}
+  out_size: ${model.out_size}
+  hid_size: 200
+
+env: "hopper--stand"
+
+env_dataset_size: 1000
+validation_ratio: 0.1
+dynamics_model_batch_size: 256
+initial_exploration_steps: 20
+num_epochs: 100
+freq_train_dyn_model: 100
+patience: 50
+rollouts_per_step: 40
+rollout_horizon: 15 # TODO replace by thresholded linear
+rollout_batch_size: 32
+sac_buffer_capacity: ???
+sac_samples_action: true
+num_sac_updates_per_rollout: 100
+
+seed: 0
+
+device: "cuda:0"
+
+log_frequency: 100
+log_save_tb: false
+
+
+experiment: test_exp
+
+hydra:
+    run:
+        dir: ./exp/mbrl/${env}/${now:%Y.%m.%d}/${now:%H%M}_${experiment}

mbrl/mbpo.py

@@ -1,12 +1,13 @@
+import os
 from typing import Callable, Tuple
 
-import dmc2gym
 import gym
+import hydra.utils
 import numpy as np
+import omegaconf
 import pytorch_sac
 import torch
 
-import mbrl.env.termination_fns as termination_fns
 import mbrl.models as models
 import mbrl.replay_buffer as replay_buffer
 
@@ -17,8 +18,9 @@ def collect_random_trajectories(
     env_dataset_test: replay_buffer.IterableReplayBuffer,
     steps_to_collect: int,
     val_ratio: float,
+    rng: np.random.RandomState,
 ):
-    indices = np.random.permutation(steps_to_collect)
+    indices = rng.permutation(steps_to_collect)
     n_train = int(steps_to_collect * (1 - val_ratio))
     indices_train = set(indices[:n_train])
 
@@ -39,109 +41,137 @@ def collect_random_trajectories(
                 return
 
 
-def rollout_model(
-    env: gym.Env,
-    model: models.Model,
+def rollout_model_and_populate_sac_buffer(
+    model_env: models.ModelEnv,
     env_dataset: replay_buffer.BootstrapReplayBuffer,
-    termination_fn: Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray],
-    obs_shape: Tuple[int],
-    act_shape: Tuple[int],
-    sac_buffer_capacity: int,
-    num_rollouts: int,
+    agent: pytorch_sac.SACAgent,
+    sac_buffer: pytorch_sac.ReplayBuffer,
+    sac_samples_action: bool,
     rollout_horizon: int,
     batch_size: int,
-    device: torch.device,
-) -> pytorch_sac.ReplayBuffer:
-    model_env = models.ModelEnv(env, model, termination_fn)
-    sac_buffer = pytorch_sac.ReplayBuffer(
-        obs_shape, act_shape, sac_buffer_capacity, device
-    )
-    for _ in range(num_rollouts):
-        initial_obs, action, *_ = env_dataset.sample(batch_size, ensemble=False)
-        obs = model_env.reset(initial_obs_batch=initial_obs)
-        for i in range(rollout_horizon):
-            pred_next_obs, pred_rewards, pred_dones, _ = model_env.step(action)
-            # TODO consider changing sac_buffer to vectorize this loop
-            for j in range(batch_size):
-                sac_buffer.add(
-                    obs[j],
-                    action[j],
-                    pred_rewards[j],
-                    pred_next_obs[j],
-                    pred_dones[j],
-                    pred_dones[j],
-                )
-            obs = pred_next_obs
+):
 
-    return sac_buffer
+    initial_obs, action, *_ = env_dataset.sample(batch_size, ensemble=False)
+    obs = model_env.reset(initial_obs_batch=initial_obs)
+    for i in range(rollout_horizon):
+        action = agent.act(obs, sample=sac_samples_action, batched=True)
+        pred_next_obs, pred_rewards, pred_dones, _ = model_env.step(action)
+        # TODO change sac_buffer to vectorize this loop (the batch size will be really large)
+        for j in range(batch_size):
+            sac_buffer.add(
+                obs[j],
+                action[j],
+                pred_rewards[j],
+                pred_next_obs[j],
+                pred_dones[j],
+                pred_dones[j],
+            )
+        obs = pred_next_obs
 
 
-def mbpo(
+def train(
     env: gym.Env,
     termination_fn: Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray],
     device: torch.device,
+    cfg: omegaconf.DictConfig,
 ):
+    # ------------------- Initialization -------------------
     obs_shape = env.observation_space.shape
     act_shape = env.action_space.shape
 
-    # PARAMS TO MOVE TO A CONFIG FILE
-    ensemble_size = 7
-    val_ratio = 0.1
-    buffer_capacity = 1000
-    batch_size = 256
-    steps_to_collect = 100
-    num_epochs = 100
-    freq_train_dyn_model = 10
-    patience = 50
-    rollouts_per_step = 40
-    rollout_horizon = 15
-    sac_buffer_capacity = 10000
-
-    # Agent
-    # agent = pytorch_sac.SACAgent()
-
-    # Creating environment datasets
+    cfg.agent.obs_dim = obs_shape[0]
+    cfg.agent.action_dim = act_shape[0]
+    cfg.agent.action_range = [
+        float(env.action_space.low.min()),
+        float(env.action_space.high.max()),
+    ]
+    agent = hydra.utils.instantiate(cfg.agent)
+
+    work_dir = os.getcwd()
+    logger = pytorch_sac.Logger(
+        work_dir, save_tb=cfg.log_save_tb, log_frequency=cfg.log_frequency, agent="sac"
+    )
+
+    rng = np.random.RandomState(cfg.seed)
+
+    # -------------- Create initial env. dataset --------------
     env_dataset_train = replay_buffer.BootstrapReplayBuffer(
-        buffer_capacity, batch_size, ensemble_size, obs_shape, act_shape
+        cfg.env_dataset_size,
+        cfg.dynamics_model_batch_size,
+        cfg.model.ensemble_size,
+        obs_shape,
+        act_shape,
     )
+    val_buffer_capacity = int(cfg.env_dataset_size * cfg.validation_ratio)
     env_dataset_val = replay_buffer.IterableReplayBuffer(
-        int(buffer_capacity * val_ratio), batch_size, obs_shape, act_shape
+        val_buffer_capacity, cfg.dynamics_model_batch_size, obs_shape, act_shape
     )
+    # TODO replace this with some exploration policy
     collect_random_trajectories(
-        env, env_dataset_train, env_dataset_val, steps_to_collect, val_ratio
+        env,
+        env_dataset_train,
+        env_dataset_val,
+        cfg.initial_exploration_steps,
+        cfg.validation_ratio,
+        rng,
     )
 
-    # Training loop
-    model_in_size = obs_shape[0] + act_shape[0]
-    model_out_size = obs_shape[0] + 1
-    ensemble = models.Ensemble(
-        models.GaussianMLP, ensemble_size, model_in_size, model_out_size, device
+    # ---------------------------------------------------------
+    # --------------------- Training Loop ---------------------
+    cfg.model.in_size = obs_shape[0] + act_shape[0]
+    cfg.model.out_size = obs_shape[0] + 1
+
+    ensemble = hydra.utils.instantiate(cfg.model)
+
+    sac_buffer_capacity = (
+        cfg.rollouts_per_step * cfg.rollout_horizon * cfg.rollout_batch_size
     )
-    for epoch in range(num_epochs):
-        if epoch % freq_train_dyn_model == 0:
-            train_loss, val_score = models.train_dyn_ensemble(
-                ensemble,
-                env_dataset_train,
-                device,
-                dataset_val=env_dataset_val,
-                patience=patience,
+
+    updates_made = 0
+    env_steps = 0
+    model_env = models.ModelEnv(env, ensemble, termination_fn)
+    for epoch in range(cfg.num_epochs):
+        obs = env.reset()
+        done = False
+        while not done:
+            # --------------- Env. Step and adding to model dataset -----------------
+            action = agent.act(obs)
+            next_obs, reward, done, _ = env.step(action)
+            if rng.random() < cfg.validation_ratio:
+                env_dataset_val.add(obs, action, next_obs, reward, done)
+            else:
+                env_dataset_train.add(obs, action, next_obs, reward, done)
+            obs = next_obs
+
+            # --------------- Model Training -----------------
+            if env_steps % cfg.freq_train_dyn_model == 0:
+                train_loss, val_score = models.train_dyn_ensemble(
+                    ensemble,
+                    env_dataset_train,
+                    device,
+                    dataset_val=env_dataset_val,
+                    patience=cfg.patience,
+                )
+
+            # --------------- Agent Training -----------------
+            sac_buffer = pytorch_sac.ReplayBuffer(
+                obs_shape, act_shape, sac_buffer_capacity, device
             )
+            for _ in range(cfg.rollouts_per_step):
+                rollout_model_and_populate_sac_buffer(
+                    model_env,
+                    env_dataset_train,
+                    agent,
+                    sac_buffer,
+                    cfg.sac_samples_action,
+                    cfg.rollout_horizon,
+                    cfg.rollout_batch_size,
+                )
+
+                for _ in range(cfg.num_sac_updates_per_rollout):
+                    agent.update(sac_buffer, logger, updates_made)
+                    updates_made += 1
+
+            logger.dump(updates_made, save=True)
 
-        sac_buffer = rollout_model(
-            env,
-            ensemble,
-            env_dataset_train,
-            termination_fn,
-            obs_shape,
-            act_shape,
-            sac_buffer_capacity,
-            rollouts_per_step,
-            rollout_horizon,
-            batch_size,
-            device,
-        )
-
-
-if __name__ == "__main__":
-    _env = dmc2gym.make(domain_name="hopper", task_name="stand")
-    mbpo(_env, termination_fns.hopper, torch.device("cuda:0"))
+            env_steps += 1

mbrl/models.py

@@ -3,7 +3,9 @@
 from typing import Tuple, List, Optional, Type, Dict, Sequence, Union
 
 import gym
+import hydra.utils
 import numpy as np
+import omegaconf
 import torch
 from torch import nn as nn, optim as optim
 from torch.nn import functional as F
@@ -87,24 +89,22 @@ def eval_score(self, model_in: torch.Tensor, target: torch.Tensor) -> float:
 class Ensemble(Model):
     def __init__(
         self,
-        cls: Type[Model],
-        num_members: int,
+        ensemble_size: int,
         in_size: int,
         out_size: int,
         device: torch.device,
-        *model_args,
+        member_cfg: omegaconf.DictConfig,
         optim_lr: float = 0.0075,
-        seed: Optional[int] = None,
-        **model_kwargs,
     ):
         super().__init__(in_size, out_size, device)
         self.members = []
         self.optimizers = []
-        for i in range(num_members):
-            model = cls(in_size, out_size, device, *model_args, **model_kwargs)
+        for i in range(ensemble_size):
+            model = hydra.utils.instantiate(member_cfg)
+            # model = member_cls(in_size, out_size, device, *model_args, **model_kwargs)
             self.members.append(model.to(device))
             self.optimizers.append(optim.Adam(model.parameters(), lr=optim_lr))
-        self.rng = np.random.RandomState(seed)
+        self.rng = np.random.RandomState()
 
     def __len__(self):
         return len(self.members)
@@ -261,7 +261,7 @@ def step(self, actions: np.ndarray):
             model_in = torch.from_numpy(
                 np.concatenate([self._current_obs, actions], axis=1)
             ).to(self.model.device)
-            model_out = self.model(model_in).cpu().numpy()[0]
+            model_out = self.model(model_in)[0].cpu().numpy()
             next_observs = model_out[:, :-1]
             rewards = model_out[:, -1]
             dones = self.termination_fn(actions, next_observs)