Retrace returns for off-policy RL.

Author: vladfi1

phillip/RL.py

@@ -121,23 +121,24 @@ def process_experiences(f, keys):
             delayed = live.copy()
             delayed.update(process_experiences(lambda t: t[:,self.config.delay:], ['state', 'reward']))
 
-          policy_args = live
-          critic_args = delayed
+          policy_args = live.copy()
+          critic_args = delayed.copy()
 
           print("Creating train ops")
 
           train_ops = []
 
-          if self.train_policy or self.train_critic:
+          if self.train_policy:
+            probs = self.policy.probs(**policy_args)
+            critic_args.update(**probs)
+            
             train_critic, targets, advantages = self.critic(**critic_args)
-          
-          if self.train_critic:
             train_ops.append(train_critic)
+            
+            probs.update(advantages=advantages)
+            train_policy = self.policy.train(**probs)
+            train_ops.append(train_policy)
 
-          if self.train_policy:
-            policy_args.update(advantages=tf.stop_gradient(advantages), targets=targets)
-            train_ops.append(self.policy.train(**policy_args))
-
           if self.train_model:
             train_ops.append(self.model.train(**delayed))
 

phillip/ac.py

@@ -50,39 +50,47 @@ def __init__(self, embedGame, embedAction, global_step, rlConfig, **kwargs):
 
     self.actor = net
 
-  def train(self, state, prev_action, action, prob, advantages, **unused):
+  def train(self, target_log_probs, advantages, **unused):
+    train_log_probs = target_log_probs[:,:-1] # last state has no advantage
+    actor_gain = tf.reduce_mean(tf.mul(train_log_probs, advantages - self.entropy_scale))
+    
+    actor_params = self.actor.getVariables()
+      
+    def metric(log_p1, log_p2):
+      return tf.reduce_mean(tf.squared_difference(log_p1, log_p2))
+    
+    return self.optimizer.optimize(-actor_gain, actor_params, target_log_probs, metric)
+  
+  def probs(self, state, prev_action, action, prob, **unused):
     embedded_state = self.embedGame(state)
     embedded_prev_action = self.embedAction(prev_action)
     history = RL.makeHistory(embedded_state, embedded_prev_action, self.rlConfig.memory)
-    
-    actor_probs = self.actor(history)
-    log_actor_probs = tf.log(actor_probs)
 
+    actions = self.embedAction(action[:,self.rlConfig.memory:])
+
+    actor_probs = self.actor(history)
+    real_actor_probs = tfl.batch_dot(actions, actor_probs)
+    
+    """
     entropy = - tfl.batch_dot(actor_probs, log_actor_probs)
     entropy_avg = tfl.power_mean(self.entropy_power, entropy)
     tf.scalar_summary('entropy_avg', entropy_avg)
     tf.scalar_summary('entropy_min', tf.reduce_min(entropy))
     tf.histogram_summary('entropy', entropy)
-
-    actions = self.embedAction(action[:,self.rlConfig.memory:])
-
-    real_actor_probs = tfl.batch_dot(actions, actor_probs)
-    prob_ratios = prob[:,self.rlConfig.memory:] / real_actor_probs
-    tf.scalar_summary('kl', tf.reduce_mean(tf.log(prob_ratios)))
-
-    real_log_actor_probs = tfl.batch_dot(actions, log_actor_probs)
-    train_log_actor_probs = real_log_actor_probs[:,:-1] # last state has no advantage
-    actor_gain = tf.reduce_mean(tf.mul(train_log_actor_probs, tf.stop_gradient(advantages)))
-    #tf.scalar_summary('actor_gain', actor_gain)
+    """
 
-    actor_loss = - (actor_gain + self.entropy_scale * entropy_avg)
+    tf.scalar_summary('entropy_avg', -tf.reduce_mean(tf.log(prob)))
 
-    actor_params = self.actor.getVariables()
-      
-    def metric(p1, p2):
-      return tf.reduce_mean(tfl.kl(p1, p2))
+    behavior_probs = prob[:,self.rlConfig.memory:]
+    ratios = real_actor_probs / behavior_probs
+    
+    tf.scalar_summary('kl', -tf.reduce_mean(tf.log(ratios)))
 
-    return self.optimizer.optimize(actor_loss, actor_params, log_actor_probs, metric)
+    return dict(
+      target_probs = real_actor_probs,
+      target_log_probs = tf.log(real_actor_probs),
+      ratios = ratios
+    )
 
   def getPolicy(self, state, **unused):
     return self.actor(state)

phillip/critic.py

@@ -10,6 +10,7 @@ class Critic(Default):
     Option('critic_learning_rate', type=float, default=1e-4),
     Option('gae_lambda', type=float, default=1., help="Generalized Advantage Estimation"),
     Option('fix_scopes', type=bool, default=False),
+    Option('retrace', type=bool, default=True, help="Retrace(lambda) - correct for off-policy behavior"),
   ]
 
   _members = [
@@ -41,7 +42,7 @@ def __init__(self, embedGame, embedAction, scope='critic', **kwargs):
 
     self.variables = self.net.getVariables()
 
-  def __call__(self, state, prev_action, reward, **unused):
+  def __call__(self, state, prev_action, reward, ratios, **unused):
     embedded_state = self.embedGame(state)
     embedded_prev_action = self.embedAction(prev_action)
     history = makeHistory(embedded_state, embedded_prev_action, self.rlConfig.memory)
@@ -53,7 +54,13 @@ def __call__(self, state, prev_action, reward, **unused):
     rewards = reward[:,self.rlConfig.memory:]
     deltaVs = rewards + self.rlConfig.discount * values[:,1:] - trainVs
 
-    advantages = tfl.discount2(deltaVs, self.rlConfig.discount * self.gae_lambda)
+    if self.retrace:
+      discounts = tf.minimum(1., ratios[:,:-1])
+    else:
+      discounts = tf.ones_like(trainVs)
+    
+    discounts *= self.rlConfig.discount * self.gae_lambda
+    advantages = tfl.discount2(deltaVs, discounts)
 
     targets = trainVs + advantages
     # targets = tfl.discount2(rewards, self.rlConfig.discount, lastV)

phillip/tf_lib.py

@@ -367,7 +367,7 @@ def discount(values, gamma, initial=None):
 
   return tf.pack(values, axis=1)
 
-def discount2(values, gamma, initial=None):
+def discount2(values, gammas, initial=None):
   """Compute returns from rewards.
   
   Uses tf.while_loop instead of unrolling in python.
@@ -382,7 +382,7 @@ def discount2(values, gamma, initial=None):
   """
 
   def body(i, prev, returns):
-    next = values[:,i] + gamma * prev
+    next = values[:,i] + gammas[:,i] * prev
     next.set_shape(prev.get_shape())
 
     returns = returns.write(i, next)