From 9c58ab62c0ce3bb5362243623e44fc08ede4c12c Mon Sep 17 00:00:00 2001
From: NM512 <morihira3513@gmail.com>
Date: Sat, 17 Jun 2023 16:59:40 +0900
Subject: [PATCH] introduced return used in author's code

---
 models.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/models.py b/models.py
index 8fadf7d..2c00597 100644
--- 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,