fix eval reward sum

pytorch · vmoens · Dec 14, 2023 · Mar 28, 2023 · Mar 28, 2023 · Mar 28, 2023
commit eb7cee021bcac313c3b376117a9da4d19d432861
diff --git a/examples/iql/iql_offline.py b/examples/iql/iql_offline.py
@@ -75,16 +75,18 @@ def main(cfg: "DictConfig"):  # noqa: F821
                 )
 
         if r0 is None:
-            r0 = eval_td["reward"].mean().item()
+            r0 = eval_td["reward"].sum(1).mean().item()
         if l0 is None:
             l0 = loss_val.item()
 
         for key, value in loss_vals.items():
             logger.log_scalar(key, value.item(), i)
-        logger.log_scalar("reward_evaluation", eval_td["reward"].mean().item(), i)
+        logger.log_scalar(
+            "reward_evaluation", eval_td["reward"].sum(1).mean().item(), i
+        )
 
         pbar.set_description(
-            f"loss: {loss_val.item(): 4.4f} (init: {l0: 4.4f}), reward: {eval_td['reward'].mean(): 4.4f} (init={r0: 4.4f})"
+            f"loss: {loss_val.item(): 4.4f} (init: {l0: 4.4f}), reward: {eval_td['reward'].sum(1).mean(): 4.4f} (init={r0: 4.4f})"
         )
 
 

diff --git a/examples/iql/utils.py b/examples/iql/utils.py
@@ -268,7 +268,7 @@ def make_offline_replay_buffer(rb_cfg, state_dict):
     )
     data.append_transform(
         ObservationNorm(
-            in_keys=["observation_vector", ("next", "observation_vector")],
+            in_keys=["observation_vector"],
             loc=state_dict["transforms.2.loc"],
             scale=state_dict["transforms.2.scale"],
             standard_normal=state_dict["transforms.2.standard_normal"],