diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 60d9671..ffe2064 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -143,7 +143,19 @@ def ppo_loss( pg_loss1 = -mb_advantages * ratio pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() - v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() + old_value = mb_returns - mb_advantages + + v_clipped = old_value + jnp.clip( + newvalue - old_value, + -args.clip_coef, + args.clip_coef + ) + + v_loss_unclipped = (newvalue - mb_returns) ** 2 + v_loss_clipped = (v_clipped - mb_returns) ** 2 + + v_loss = 0.5 * jnp.maximum(v_loss_unclipped, v_loss_clipped).mean() + entropy_loss = entropy.mean() loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))