refactor(reforcexy): make n_stack discovery independant of the # of env

		     wrappers

Signed-off-by: Jérôme Benoit <jerome.benoit@piment-noir.org>

diff --git a/ReforceXY/user_data/freqaimodels/ReforceXY.py b/ReforceXY/user_data/freqaimodels/ReforceXY.py
index c12bad70e8a296a759e08f779394a23399fc7b07..10252681ce4b1a521ea333fa5d6cd2687f21fcfc 100644 (file)
--- a/ReforceXY/user_data/freqaimodels/ReforceXY.py
+++ b/ReforceXY/user_data/freqaimodels/ReforceXY.py
@@ -1485,13 +1485,15 @@ class InfoMetricsCallback(TensorboardCallback):
             float(_lr) if isinstance(_lr, (int, float, np.floating)) else "lr_schedule"
         )
         n_stack = 1
-        if self.training_env is not None and hasattr(self.training_env, "get_attr"):
-            try:
-                stacks = self.training_env.get_attr("n_stack")
-                if isinstance(stacks, (list, tuple)) and stacks and stacks[0]:
-                    n_stack = int(stacks[0])
-            except Exception:
-                pass
+        training_env = getattr(self, "training_env", None)
+        while training_env is not None:
+            if hasattr(training_env, "n_stack"):
+                try:
+                    n_stack = int(getattr(training_env, "n_stack"))
+                except Exception:
+                    pass
+                break
+            training_env = getattr(training_env, "venv", None)
         hparam_dict: Dict[str, Any] = {
             "algorithm": self.model.__class__.__name__,
             "n_envs": int(self.model.n_envs),