SHAC: add target network

Author: peabody124

brax/training/agents/shac/losses.py

@@ -83,29 +83,32 @@ def compute_shac_policy_loss(
   truncation = data.extras['state_extras']['truncation']
   termination = (1 - data.discount) * (1 - truncation)
 
-  horizon = rewards.shape[0]
+  # Append terminal values to get [v1, ..., v_t+1]
+  values_t_plus_1 = jnp.concatenate(
+      [values[1:], jnp.expand_dims(terminal_values, 0)], axis=0)
 
+  # jax implementation of https://github.com/NVlabs/DiffRL/blob/a4c0dd1696d3c3b885ce85a3cb64370b580cb913/algorithms/shac.py#L227
   def sum_step(carry, target_t):
-    gam, acc = carry
-    reward, v, truncation,  termination = target_t
-    acc = acc + jnp.where(truncation + termination, gam * v, gam * reward)
+    gam, rew_acc = carry
+    reward, v, termination = target_t
+
+    # clean up gamma and rew_acc for done envs, otherwise update
+    rew_acc = jnp.where(termination, 0, rew_acc + gam * reward)
     gam = jnp.where(termination, 1.0, gam * discounting)
-    return (gam, acc), (acc)
 
-  acc = terminal_values * (discounting ** horizon) * (1-termination[-1]) * (1-truncation[-1])
-  jax.debug.print('acc shape: {x}', x=acc.shape)
-  gam = jnp.ones_like(terminal_values)
-  (_, acc), (temp) = jax.lax.scan(sum_step, (gam, acc),
-      (rewards, values, truncation, termination))
+    return (gam, rew_acc), (gam, rew_acc)
 
-  policy_loss = -jnp.mean(acc) / horizon
+  rew_acc = jnp.zeros_like(terminal_values)
+  gam = jnp.ones_like(terminal_values)
+  (gam, last_rew_acc), (gam_acc, rew_acc) = jax.lax.scan(sum_step, (gam, rew_acc),
+      (rewards, values, termination))
 
-  # inspect the data for one of the rollouts
-  jax.debug.print('obs={o}, obs_next={n}, values={v}, reward={r}, truncation={t}, terminal={s}',
-      v=values[:, 0], o=data.observation[:,0], r=data.reward[:,0],
-      t=truncation[:, 0], s=termination[:,0], n=data.next_observation[:, 0])
+  policy_loss = jnp.sum(-last_rew_acc - gam * terminal_values)
+  # for trials that are truncated (i.e. hit the episode length) include reward for
+  # terminal state. otherwise, the trial was aborted and should receive zero additional
+  policy_loss = policy_loss + jnp.sum((-rew_acc - gam_acc * jnp.where(truncation, values_t_plus_1, 0)) * termination)
+  policy_loss = policy_loss / values.shape[0] / values.shape[1]
 
-  jax.debug.print('loss={l}, r={r}', l=policy_loss, r=temp[:,0])
 
   # Entropy reward
   policy_logits = policy_apply(normalizer_params, policy_params,
@@ -122,68 +125,6 @@ def sum_step(carry, target_t):
   }
 
 
-
-def compute_target_values(truncation: jnp.ndarray,
-                          termination: jnp.ndarray,
-                          rewards: jnp.ndarray,
-                          values: jnp.ndarray,
-                          bootstrap_value: jnp.ndarray,
-                          discount: float = 0.99,
-                          lambda_: float = 0.95,
-                          td_lambda=True):
-  """Calculates the target values.
-
-  This implements Eq. 7 of 2204.07137
-  https://github.com/NVlabs/DiffRL/blob/main/algorithms/shac.py#L349
-
-  Args:
-    truncation: A float32 tensor of shape [T, B] with truncation signal.
-    termination: A float32 tensor of shape [T, B] with termination signal.
-    rewards: A float32 tensor of shape [T, B] containing rewards generated by
-      following the behaviour policy.
-    values: A float32 tensor of shape [T, B] with the value function estimates
-      wrt. the target policy.
-    bootstrap_value: A float32 of shape [B] with the value function estimate at
-      time T.
-    discount: TD discount.
-
-  Returns:
-    A float32 tensor of shape [T, B].
-  """
-  truncation_mask = 1 - truncation
-  # Append bootstrapped value to get [v1, ..., v_t+1]
-  values_t_plus_1 = jnp.concatenate(
-      [values[1:], jnp.expand_dims(bootstrap_value, 0)], axis=0)
-
-  if td_lambda:
-
-    def compute_v_st(carry, target_t):
-      Ai, Bi, lam = carry
-      reward, truncation_mask, vtp1, termination = target_t
-      # TODO: should figure out how to handle termination
-
-      lam = lam * lambda_ * (1 - termination) + termination
-      Ai = (1 - termination) * (lam * discount * Ai + discount * vtp1 + (1. - lam) / (1. - lambda_) * reward)
-      Bi = discount * (vtp1 * termination + Bi * (1.0 - termination)) + reward
-      vs = (1.0 - lambda_) * Ai + lam * Bi
-
-      return (Ai, Bi, lam), (vs)
-
-    Ai = jnp.ones_like(bootstrap_value)
-    Bi = jnp.zeros_like(bootstrap_value)
-    lam = jnp.ones_like(bootstrap_value)
-
-    (_, _, _), (vs) = jax.lax.scan(compute_v_st, (Ai, Bi, lam),
-        (rewards, truncation_mask, values_t_plus_1, termination),
-        length=int(truncation_mask.shape[0]),
-        reverse=True)
-
-  else:
-    vs = rewards + discount * values_t_plus_1
-
-  return jax.lax.stop_gradient(vs)
-
-
 def compute_shac_critic_loss(
     params: Params,
     normalizer_params: Any,
@@ -192,9 +133,13 @@ def compute_shac_critic_loss(
     shac_network: shac_networks.SHACNetworks,
     discounting: float = 0.9,
     reward_scaling: float = 1.0,
-    lambda_: float = 0.95) -> Tuple[jnp.ndarray, types.Metrics]:
+    lambda_: float = 0.95,
+    td_lambda: bool = True) -> Tuple[jnp.ndarray, types.Metrics]:
   """Computes SHAC critic loss.
 
+  This implements Eq. 7 of 2204.07137
+  https://github.com/NVlabs/DiffRL/blob/main/algorithms/shac.py#L349
+
   Args:
     params: Value network parameters,
     normalizer_params: Parameters of the normalizer.
@@ -207,8 +152,7 @@ def compute_shac_critic_loss(
     discounting: discounting,
     reward_scaling: reward multiplier.
     lambda_: Lambda for TD value updates
-    clipping_epsilon: Policy loss clipping epsilon
-    normalize_advantage: whether to normalize advantage estimate
+    td_lambda: whether to use a TD-Lambda value target
 
   Returns:
     A tuple (loss, metrics)
@@ -218,25 +162,47 @@ def compute_shac_critic_loss(
 
   data = jax.tree_util.tree_map(lambda x: jnp.swapaxes(x, 0, 1), data)
 
-  baseline = value_apply(normalizer_params, params, data.observation)
-  bootstrap_value = value_apply(normalizer_params, params, data.next_observation[-1])
+  values = value_apply(normalizer_params, params, data.observation)
+  terminal_value = value_apply(normalizer_params, params, data.next_observation[-1])
 
   rewards = data.reward * reward_scaling
   truncation = data.extras['state_extras']['truncation']
   termination = (1 - data.discount) * (1 - truncation)
 
-  vs = compute_target_values(
-      truncation=truncation,
-      termination=termination,
-      rewards=rewards,
-      values=baseline,
-      bootstrap_value=bootstrap_value,
-      discount=discounting,
-      lambda_=lambda_)
+  # Append terminal values to get [v1, ..., v_t+1]
+  values_t_plus_1 = jnp.concatenate(
+      [values[1:], jnp.expand_dims(terminal_value, 0)], axis=0)
+
+  # compute target values
+  if td_lambda:
+
+    def compute_v_st(carry, target_t):
+      Ai, Bi, lam = carry
+      reward, vtp1, termination = target_t
+
+      reward = reward * termination
+
+      lam = lam * lambda_ * (1 - termination) + termination
+      Ai = (1 - termination) * (lam * discounting * Ai + discounting * vtp1 + (1. - lam) / (1. - lambda_) * reward)
+      Bi = discounting * (vtp1 * termination + Bi * (1.0 - termination)) + reward
+      vs = (1.0 - lambda_) * Ai + lam * Bi
+
+      return (Ai, Bi, lam), (vs)
+
+    Ai = jnp.ones_like(terminal_value)
+    Bi = jnp.zeros_like(terminal_value)
+    lam = jnp.ones_like(terminal_value)
+    (_, _, _), (vs) = jax.lax.scan(compute_v_st, (Ai, Bi, lam),
+        (rewards, values_t_plus_1, termination),
+        length=int(termination.shape[0]),
+        reverse=True)
+
+  else:
+    vs = rewards + discounting * values_t_plus_1
 
-  v_error = vs - baseline
-  v_loss = jnp.mean(v_error * v_error) * 0.5 * 0.5
+  target_values = jax.lax.stop_gradient(vs)
 
+  v_loss = jnp.mean((target_values - values) ** 2)
 
   total_loss = v_loss
   return total_loss, {

brax/training/agents/shac/train.py

@@ -53,6 +53,7 @@ class TrainingState:
   policy_params: Params
   value_optimizer_state: optax.OptState
   value_params: Params
+  target_value_params: Params
   normalizer_params: running_statistics.RunningStatisticsState
   env_steps: jnp.ndarray
 
@@ -80,6 +81,7 @@ def train(environment: envs.Env,
           num_evals: int = 1,
           normalize_observations: bool = False,
           reward_scaling: float = 1.,
+          tau: float = 0.005,  # this is 1-alpha from the original paper
           lambda_: float = .95,
           deterministic_eval: bool = False,
           network_factory: types.NetworkFactory[
@@ -222,7 +224,7 @@ def training_step(
     key_sgd, key_generate_unroll, new_key = jax.random.split(key, 3)
 
     (policy_loss, (state, data, policy_metrics)), policy_params, policy_optimizer_state = policy_gradient_update_fn(
-        training_state.policy_params, training_state.value_params,
+        training_state.policy_params, training_state.target_value_params,
         training_state.normalizer_params, state, key_generate_unroll,
         optimizer_state=training_state.policy_optimizer_state)
 
@@ -238,13 +240,18 @@ def training_step(
         (training_state.value_optimizer_state, training_state.value_params, key_sgd), (),
         length=num_updates_per_batch)
 
+    target_value_params = jax.tree_util.tree_map(
+        lambda x, y: x * (1 - tau) + y * tau, training_state.target_value_params,
+        value_params)
+
     metrics.update(policy_metrics)
 
     new_training_state = TrainingState(
         policy_optimizer_state=policy_optimizer_state,
         policy_params=policy_params,
         value_optimizer_state=value_optimizer_state,
         value_params=value_params,
+        target_value_params=target_value_params,
         normalizer_params=training_state.normalizer_params,
         env_steps=training_state.env_steps + env_step_per_training_step)
     return (new_training_state, state, new_key), metrics
@@ -298,6 +305,7 @@ def training_epoch_with_timing(
       policy_params=policy_init_params,
       value_optimizer_state=value_optimizer.init(value_init_params),
       value_params=value_init_params,
+      target_value_params=value_init_params,
       normalizer_params=running_statistics.init_state(
           specs.Array((env.observation_size,), jnp.float32)),
       env_steps=0)
@@ -329,7 +337,7 @@ def training_epoch_with_timing(
   if process_id == 0 and num_evals > 1:
     metrics = evaluator.run_evaluation(
         _unpmap(
-            (training_state.normalizer_params, training_state.params.policy)),
+            (training_state.normalizer_params, training_state.policy_params)),
         training_metrics={})
     logging.info(metrics)
     progress_fn(0, metrics)