updates to pcl_rl according to most recent version of Trust-PCL paper

Author: ofirnachum

research/pcl_rl/README.md

@@ -67,20 +67,27 @@ python trainer.py --logtostderr --batch_size=25 --env=HalfCheetah-v1 \
   --max_divergence=0.05 --value_opt=best_fit --critic_weight=0.0 \
 ```
 
-Run Mujoco task with Trust-PCL:
+To run Mujoco task using Trust-PCL (off-policy) use the below command.
+It should work well across all environments, given that you
+search sufficiently among
+
+(1) max_divergence (0.001, 0.0005, 0.002 are good values),
+
+(2) rollout (1, 5, 10 are good values),
+
+(3) tf_seed (need to average over enough random seeds).
 
 ```
 python trainer.py --logtostderr --batch_size=1 --env=HalfCheetah-v1 \
-  --validation_frequency=50 --rollout=10 --critic_weight=0.0 \
-  --gamma=0.995 --clip_norm=40 --learning_rate=0.002 \
-  --replay_buffer_freq=1 --replay_buffer_size=20000 \
-  --replay_buffer_alpha=0.1 --norecurrent --objective=pcl \
-  --max_step=100 --tau=0.0 --eviction=fifo --max_divergence=0.001 \
-  --internal_dim=64 --cutoff_agent=1000 \
-  --replay_batch_size=25 --nouse_online_batch --batch_by_steps \
-  --sample_from=target --value_opt=grad --value_hidden_layers=2 \
-  --update_eps_lambda --unify_episodes --clip_adv=1.0 \
-  --target_network_lag=0.99 --prioritize_by=step
+  --validation_frequency=250 --rollout=1 --critic_weight=1.0 --gamma=0.995 \
+  --clip_norm=40 --learning_rate=0.0001 --replay_buffer_freq=1 \
+  --replay_buffer_size=5000 --replay_buffer_alpha=0.001 --norecurrent \
+  --objective=pcl --max_step=10 --cutoff_agent=1000 --tau=0.0 --eviction=fifo \
+  --max_divergence=0.001 --internal_dim=256 --replay_batch_size=64 \
+  --nouse_online_batch --batch_by_steps --value_hidden_layers=2 \
+  --update_eps_lambda --nounify_episodes --target_network_lag=0.99 \
+  --sample_from=online --clip_adv=1 --prioritize_by=step --num_steps=1000000 \
+  --noinput_prev_actions --use_target_values --tf_seed=57
 ```
 
 Run Mujoco task with PCL constraint trust region:

research/pcl_rl/controller.py

@@ -109,13 +109,14 @@ def __init__(self, env, env_spec, internal_dim,
     self.episode_running_rewards = np.zeros(len(self.env))
     self.episode_running_lengths = np.zeros(len(self.env))
     self.episode_rewards = []
+    self.greedy_episode_rewards = []
     self.episode_lengths = []
     self.total_rewards = []
 
     self.best_batch_rewards = None
 
-  def setup(self):
-    self.model.setup()
+  def setup(self, train=True):
+    self.model.setup(train=train)
 
   def initial_internal_state(self):
     return np.zeros(self.model.policy.rnn_state_dim)
@@ -187,7 +188,7 @@ def _sample_episodes(self, sess, greedy=False):
 
     return initial_state, all_obs, all_act, rewards, all_pad
 
-  def sample_episodes(self, sess):
+  def sample_episodes(self, sess, greedy=False):
     """Sample steps from the environment until we have enough for a batch."""
 
     # check if last batch ended with episode that was not terminated
@@ -200,7 +201,7 @@ def sample_episodes(self, sess):
     while total_steps < self.max_step * len(self.env):
       (initial_state,
        observations, actions, rewards,
-       pads) = self._sample_episodes(sess)
+       pads) = self._sample_episodes(sess, greedy=greedy)
 
       observations = zip(*observations)
       actions = zip(*actions)
@@ -249,19 +250,26 @@ def _train(self, sess,
              observations, initial_state, actions,
              rewards, terminated, pads):
     """Train model using batch."""
+    avg_episode_reward = np.mean(self.episode_rewards)
+    greedy_episode_reward = (np.mean(self.greedy_episode_rewards)
+                             if self.greedy_episode_rewards else
+                             avg_episode_reward)
+    loss, summary = None, None
     if self.use_trust_region:
       # use trust region to optimize policy
       loss, _, summary = self.model.trust_region_step(
           sess,
           observations, initial_state, actions,
           rewards, terminated, pads,
-          avg_episode_reward=np.mean(self.episode_rewards))
+          avg_episode_reward=avg_episode_reward,
+          greedy_episode_reward=greedy_episode_reward)
     else:  # otherwise use simple gradient descent on policy
       loss, _, summary = self.model.train_step(
           sess,
           observations, initial_state, actions,
           rewards, terminated, pads,
-          avg_episode_reward=np.mean(self.episode_rewards))
+          avg_episode_reward=avg_episode_reward,
+          greedy_episode_reward=greedy_episode_reward)
 
     if self.use_value_opt:  # optionally perform specific value optimization
       self.model.fit_values(
@@ -305,7 +313,8 @@ def train(self, sess):
     if self.update_eps_lambda:
       episode_rewards = np.array(self.episode_rewards)
       episode_lengths = np.array(self.episode_lengths)
-      eps_lambda = find_best_eps_lambda(episode_rewards, episode_lengths)
+      eps_lambda = find_best_eps_lambda(
+          episode_rewards[-20:], episode_lengths[-20:])
       sess.run(self.model.objective.assign_eps_lambda,
                feed_dict={self.model.objective.new_eps_lambda: eps_lambda})
 
@@ -328,10 +337,10 @@ def eval(self, sess):
     """Use greedy sampling."""
     (initial_state,
      observations, actions, rewards,
-     pads) = self._sample_episodes(sess, greedy=True)
+     pads, terminated) = self.sample_episodes(sess, greedy=True)
 
     total_rewards = np.sum(np.array(rewards) * (1 - np.array(pads)), axis=0)
-    return np.mean(total_rewards)
+    return total_rewards, self.episode_rewards
 
   def convert_from_batched_episodes(
       self, initial_state, observations, actions, rewards,
@@ -351,7 +360,7 @@ def convert_from_batched_episodes(
     for i in xrange(num_episodes):
       length = total_length[i]
       ep_initial = initial_state[i]
-      ep_obs = [obs[:length, i, ...] for obs in observations]
+      ep_obs = [obs[:length + 1, i, ...] for obs in observations]
       ep_act = [act[:length + 1, i, ...] for act in actions]
       ep_rewards = rewards[:length, i]
 

research/pcl_rl/full_episode_objective.py

@@ -42,7 +42,8 @@ def get_bonus(self, total_rewards, total_log_probs):
 
   def get(self, rewards, pads, values, final_values,
           log_probs, prev_log_probs, target_log_probs,
-          entropies, logits):
+          entropies, logits,
+          target_values, final_target_values):
     seq_length = tf.shape(rewards)[0]
 
     not_pad = tf.reshape(1 - pads, [seq_length, -1, self.num_samples])

research/pcl_rl/model.py

@@ -57,6 +57,8 @@ def setup_placeholders(self):
     # summary placeholder
     self.avg_episode_reward = tf.placeholder(
         tf.float32, [], 'avg_episode_reward')
+    self.greedy_episode_reward = tf.placeholder(
+        tf.float32, [], 'greedy_episode_reward')
 
     # sampling placeholders
     self.internal_state = tf.placeholder(tf.float32,
@@ -118,12 +120,13 @@ def setup_placeholders(self):
     self.prev_log_probs = tf.placeholder(tf.float32, [None, None],
                                          'prev_log_probs')
 
-  def setup(self):
+  def setup(self, train=True):
     """Setup Tensorflow Graph."""
 
     self.setup_placeholders()
 
     tf.summary.scalar('avg_episode_reward', self.avg_episode_reward)
+    tf.summary.scalar('greedy_episode_reward', self.greedy_episode_reward)
 
     with tf.variable_scope('model', reuse=None):
       # policy network
@@ -174,45 +177,46 @@ def setup(self):
           target_p.assign(aa * target_p + (1 - aa) * online_p)
           for online_p, target_p in zip(online_vars, target_vars)])
 
-      # evaluate objective
-      (self.loss, self.raw_loss, self.regression_target,
-       self.gradient_ops, self.summary) = self.objective.get(
-          self.rewards, self.pads,
-          self.values[:-1, :],
-          self.values[-1, :] * (1 - self.terminated),
-          self.log_probs, self.prev_log_probs, self.target_log_probs,
-          self.entropies,
-          self.logits)
-
-      self.regression_target = tf.reshape(self.regression_target, [-1])
-
-      self.policy_vars = [
-          v for v in tf.trainable_variables()
-          if '/policy_net' in v.name]
-      self.value_vars = [
-          v for v in tf.trainable_variables()
-          if '/value_net' in v.name]
-
-    # trust region optimizer
-    if self.trust_region_policy_opt is not None:
-      with tf.variable_scope('trust_region_policy', reuse=None):
-        avg_self_kl = (
-            tf.reduce_sum(sum(self.self_kls) * (1 - self.pads)) /
-            tf.reduce_sum(1 - self.pads))
-
-        self.trust_region_policy_opt.setup(
-            self.policy_vars, self.raw_loss, avg_self_kl,
-            self.avg_kl)
-
-    # value optimizer
-    if self.value_opt is not None:
-      with tf.variable_scope('trust_region_value', reuse=None):
-        self.value_opt.setup(
-            self.value_vars,
-            tf.reshape(self.values[:-1, :], [-1]),
-            self.regression_target,
-            tf.reshape(self.pads, [-1]),
-            self.regression_input, self.regression_weight)
+      if train:
+        # evaluate objective
+        (self.loss, self.raw_loss, self.regression_target,
+         self.gradient_ops, self.summary) = self.objective.get(
+            self.rewards, self.pads,
+            self.values[:-1, :],
+            self.values[-1, :] * (1 - self.terminated),
+            self.log_probs, self.prev_log_probs, self.target_log_probs,
+            self.entropies, self.logits, self.target_values[:-1, :],
+            self.target_values[-1, :] * (1 - self.terminated))
+
+        self.regression_target = tf.reshape(self.regression_target, [-1])
+
+        self.policy_vars = [
+            v for v in tf.trainable_variables()
+            if '/policy_net' in v.name]
+        self.value_vars = [
+            v for v in tf.trainable_variables()
+            if '/value_net' in v.name]
+
+        # trust region optimizer
+        if self.trust_region_policy_opt is not None:
+          with tf.variable_scope('trust_region_policy', reuse=None):
+            avg_self_kl = (
+                tf.reduce_sum(sum(self.self_kls) * (1 - self.pads)) /
+                tf.reduce_sum(1 - self.pads))
+
+            self.trust_region_policy_opt.setup(
+                self.policy_vars, self.raw_loss, avg_self_kl,
+                self.avg_kl)
+
+        # value optimizer
+        if self.value_opt is not None:
+          with tf.variable_scope('trust_region_value', reuse=None):
+            self.value_opt.setup(
+                self.value_vars,
+                tf.reshape(self.values[:-1, :], [-1]),
+                self.regression_target,
+                tf.reshape(self.pads, [-1]),
+                self.regression_input, self.regression_weight)
 
     # we re-use variables for the sampling operations
     with tf.variable_scope('model', reuse=True):
@@ -249,32 +253,42 @@ def sample_step(self, sess,
   def train_step(self, sess,
                  observations, internal_state, actions,
                  rewards, terminated, pads,
-                 avg_episode_reward=0):
+                 avg_episode_reward=0, greedy_episode_reward=0):
     """Train network using standard gradient descent."""
     outputs = [self.raw_loss, self.gradient_ops, self.summary]
     feed_dict = {self.internal_state: internal_state,
                  self.rewards: rewards,
                  self.terminated: terminated,
                  self.pads: pads,
-                 self.avg_episode_reward: avg_episode_reward}
+                 self.avg_episode_reward: avg_episode_reward,
+                 self.greedy_episode_reward: greedy_episode_reward}
+    time_len = None
     for action_place, action in zip(self.actions, actions):
+      if time_len is None:
+        time_len = len(action)
+      assert time_len == len(action)
       feed_dict[action_place] = action
     for obs_place, obs in zip(self.observations, observations):
+      assert time_len == len(obs)
       feed_dict[obs_place] = obs
 
+    assert len(rewards) == time_len - 1
+
     return sess.run(outputs, feed_dict=feed_dict)
 
 
   def trust_region_step(self, sess,
                         observations, internal_state, actions,
                         rewards, terminated, pads,
-                        avg_episode_reward=0):
+                        avg_episode_reward=0,
+                        greedy_episode_reward=0):
     """Train policy using trust region step."""
     feed_dict = {self.internal_state: internal_state,
                  self.rewards: rewards,
                  self.terminated: terminated,
                  self.pads: pads,
-                 self.avg_episode_reward: avg_episode_reward}
+                 self.avg_episode_reward: avg_episode_reward,
+                 self.greedy_episode_reward: greedy_episode_reward}
     for action_place, action in zip(self.actions, actions):
       feed_dict[action_place] = action
     for obs_place, obs in zip(self.observations, observations):