code: def value_opt_step(i,
                     opt_state,
                     opt_update,
                     value_net_apply,
                     padded_observations,
                     padded_rewards,
                     reward_mask,
                     gamma=0.99):
    
    value_params = trax_opt.get_params(opt_state)
    # Note this partial application here and argnums above in ppo_opt_step.
    g = grad(functools.partial(value_loss, value_net_apply))(
        value_params,
        padded_observations,
        padded_rewards,
        reward_mask,
        gamma=gamma)
    return opt_update(i, g, opt_state)
docs: Value optimizer step.