fix extract rnn modules

Acellera · Jun 5, 2024 · 395d06e · 395d06e
1 parent 61e6928
commit 395d06e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/scripts/hill_climb/hill_climb.py b/scripts/hill_climb/hill_climb.py
@@ -179,7 +179,8 @@ def create_env_fn():
         env = SMILESEnv(**env_kwargs)
         env = TransformedEnv(env)
         env.append_transform(InitTracker())
-        env.append_transform(get_primers_from_module(actor_training))
+        if primers := get_primers_from_module(actor_inference):
+            env.append_transform(primers)
         return env
 
     env = create_env_fn()

diff --git a/scripts/sac/pretrain_sac.py b/scripts/sac/pretrain_sac.py
@@ -345,8 +345,9 @@ def create_env_fn():
                     "train/reward": episode_rewards.mean().item(),
                     "train/min_reward": episode_rewards.min().item(),
                     "train/max_reward": episode_rewards.max().item(),
-                    "train/episode_length": episode_length.sum().item()
-                    / len(episode_length),
+                    "train/episode_length": episode_length.sum().item() / len(
+                        episode_length
+                    ),
                 }
             )
             if logger:

diff --git a/scripts/sac/sac.py b/scripts/sac/sac.py
@@ -345,8 +345,9 @@ def create_env_fn():
                     "train/reward": episode_rewards.mean().item(),
                     "train/min_reward": episode_rewards.min().item(),
                     "train/max_reward": episode_rewards.max().item(),
-                    "train/episode_length": episode_length.sum().item()
-                    / len(episode_length),
+                    "train/episode_length": episode_length.sum().item() / len(
+                        episode_length
+                    ),
                 }
             )
             if logger: