code: 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)
docs: PPO optimizer step.