def ppo_opt_step(i,
                 opt_state,
                 ppo_opt_update,
                 policy_net_apply,
                 old_policy_params,
                 value_net_apply,
                 value_net_params,
                 padded_observations,
                 padded_actions,
                 padded_rewards,
                 reward_mask,
                 gamma=0.99,
                 lambda_=0.95,
                 epsilon=0.1):
  """PPO optimizer step."""
  new_policy_params = trax_opt.get_params(opt_state)
  g = grad(
      ppo_loss, argnums=1)(
          policy_net_apply,
          new_policy_params,
          old_policy_params,
          value_net_apply,
          value_net_params,
          padded_observations,
          padded_actions,
          padded_rewards,
          reward_mask,
          gamma=gamma,
          lambda_=lambda_,
          epsilon=epsilon)
  return ppo_opt_update(i, g, opt_state)
def ppo_opt_step ( i , opt_state , ppo_opt_update , policy_net_apply , old_policy_params , value_net_apply , value_net_params , padded_observations , padded_actions , padded_rewards , reward_mask , gamma = 0.99 , lambda_ = 0.95 , epsilon = 0.1 ) : new_policy_params = trax_opt . get_params ( opt_state ) g = grad ( ppo_loss , argnums = 1 ) ( policy_net_apply , new_policy_params , old_policy_params , value_net_apply , value_net_params , padded_observations , padded_actions , padded_rewards , reward_mask , gamma = gamma , lambda_ = lambda_ , epsilon = epsilon ) return ppo_opt_update ( i , g , opt_state )
PPO optimizer step.