introduced return used in author's code

2023-06-17 16:59:40 +09:00 · 2023-06-17 16:59:40 +09:00 · 9c58ab62c0
commit 9c58ab62c0
parent 784a69f80c
1 changed files with 2 additions and 2 deletions
--- a/models.py
+++ b/models.py
@ -400,9 +400,9 @@ class ImagBehavior(nn.Module):
            reward += self._config.actor_state_entropy() * state_ent
        value = self.value(imag_feat).mode()
        target = tools.lambda_return(
-            reward[:-1],
+            reward[1:],
            value[:-1],
-            discount[:-1],
+            discount[1:],
            bootstrap=value[-1],
            lambda_=self._config.discount_lambda,
            axis=0,