Update

[ghstack-poisoned]

Author: vmoens

Diff for: torchrl/envs/custom/llm.py

@@ -325,7 +325,15 @@ def _make_next_obs(
             if self.attention_key is not None:
                 attention_mask = tensordict.get(self.attention_key)
                 n = action.shape[-1] - attention_mask.shape[-1]
-                attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[:-1] + (n,))], -1)
+                if n > 0:
+                    # It can happen that there's only one action (eg rand_action)
+                    attention_mask = torch.cat(
+                        [
+                            attention_mask,
+                            attention_mask.new_ones(attention_mask.shape[:-1] + (n,)),
+                        ],
+                        -1,
+                    )
                 nex_td.set(self.attention_key, attention_mask)
             return nex_td
 

Diff for: torchrl/envs/transforms/transforms.py

@@ -6361,10 +6361,13 @@ def _reset_func(
 
     def __repr__(self) -> str:
         class_name = self.__class__.__name__
-        default_value = {
-            key: value if isinstance(value, float) else "Callable"
-            for key, value in self.default_value.items()
-        }
+        if callable(self.default_value):
+            default_value = self.default_value
+        else:
+            default_value = {
+                key: value if isinstance(value, float) else "Callable"
+                for key, value in self.default_value.items()
+            }
         return f"{class_name}(primers={self.primers}, default_value={default_value}, random={self.random})"