From a8456e95bcbda66f80148a561b37a52d7f4ddb6b Mon Sep 17 00:00:00 2001
From: Unai Ruiz <68402420+Ruicky8@users.noreply.github.com>
Date: Tue, 3 Mar 2026 16:34:52 +0100
Subject: [PATCH] Fix missing advantage computation when reward_EMA is disabled

This PR fixes an issue where reward_EMA=False caused adv to be undefined in _compute_actor_loss.

Previously, adv was only computed inside the reward_EMA branch, which resulted in a runtime error when the option was disabled.
---
 models.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/models.py b/models.py
index ae923e8..91f756f 100644
--- a/models.py
+++ b/models.py
@@ -409,6 +409,8 @@ class ImagBehavior(nn.Module):
             metrics.update(tools.tensorstats(normed_target, "normed_target"))
             metrics["EMA_005"] = to_np(self.ema_vals[0])
             metrics["EMA_095"] = to_np(self.ema_vals[1])
+        else:
+            adv = target - base
 
         if self._config.imag_gradient == "dynamics":
             actor_target = adv