Update

vmoens · vmoens · commit a6d2c96de548 · 2025-04-17T13:53:42.000+01:00
[ghstack-poisoned]
diff --git a/torchrl/envs/transforms/llm.py b/torchrl/envs/transforms/llm.py
@@ -765,7 +765,7 @@ def _step(
         kl = curr_log_prob - log_prob
         if reward is None:
             reward = 0
-        next_tensordict.set(self.out_keys[0], reward + self.coef * kl)
+        next_tensordict.set(self.out_keys[0], reward - self.coef * kl)
         return next_tensordict
 
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
