Merge pull request #56 from yhs0602/patch-1

Bug fix for NGU module's AttributeError
RLE-Foundation · Jan 4, 2025 · 0bd71ec · 0bd71ec
2 parents eeefded + 18183b2
commit 0bd71ec
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rllte/xplore/reward/ngu.py b/rllte/xplore/reward/ngu.py
@@ -121,6 +121,8 @@ def __init__(
             weight_init=weight_init,
         )
 
+        self.mrs = mrs
+
         super().__init__(*[rnd, pseudo_counts])
 
     def compute(self, samples: Dict[str, th.Tensor], sync: bool = True) -> th.Tensor:
@@ -145,4 +147,4 @@ def compute(self, samples: Dict[str, th.Tensor], sync: bool = True) -> th.Tensor
             lifelong_rewards, th.ones_like(lifelong_rewards) * self.mrs
         )
 
-        return lifelong_rewards * episodic_rewards
+        return lifelong_rewards * episodic_rewards