No public description

PiperOrigin-RevId: 717530885
google · Jan 20, 2025 · f9083bb · f9083bb
1 parent 2415c13
commit f9083bb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/learned_optimization/research/discovered_policy_optimisation/truncated_step.py b/learned_optimization/research/discovered_policy_optimisation/truncated_step.py
@@ -236,7 +236,7 @@ def compute_drift_loss(
   if ppo_init:
     eps = 0.2
     drift_ppo = flax.linen.relu(
-        (rho_s - jnp.clip(rho_s, a_min=1 - eps, a_max=1 + eps)) * advantages)
+        (rho_s - jnp.clip(rho_s, min=1 - eps, max=1 + eps)) * advantages)
     drift = out_fn(drift_ppo + drift - 0.0001)
   else:
     drift = out_fn(drift - 0.0001)