add save function in pipeline

2024-06-16 21:47:53 +08:00
parent b9d6482d11
commit fb2ae5d2fa
10 changed files with 94 additions and 164 deletions
--- a/tensorneat/problem/rl_env/rl_jit.py
+++ b/tensorneat/problem/rl_env/rl_jit.py
@@ -43,7 +43,7 @@ class RLEnv(BaseProblem):
            assert sample_episodes > 0, "sample_size must be greater than 0"
            self.sample_policy = sample_policy
            self.sample_episodes = sample_episodes
-            self.obs_normalization = obs_normalization
+        self.obs_normalization = obs_normalization

    def setup(self, state=State()):
        if self.obs_normalization: