internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: pytorch/rl

Issue:
[BUG] Default value of the `tanh_loc` parameter of `TruncatedNormal` does not match the documentation.
## Describe the bug

The [doc string](https://github.com/pytorch/rl/blame/f3e9a1d03909d797c06f4db33635f38dc7f7efff/torchrl/modules/distributions/continuous.py#L307) of `TruncatedNormal` says that the default value of `tanh_loc` is `True`, however in the [code](https://github.com/pytorch/rl/blame/f3e9a1d03909d797c06f4db33635f38dc7f7efff/torchrl/modules/distributions/continuous.py#L325), it is set to `False`.

## Expected behavior

The values should match.

## Reason and Possible fixes

Very recently, #1192 changed the default value of the `tanh_loc` parameter of `TruncatedNormal` from `True` to `False`, but did not update the doc string of the class.

I'm not sure which value should be the default value.
The doc string seems to argue in favor of `False` to avoid numerical instability.

However, it would be good to be consistent with `IndependentNormal` and `TruncatedNormal` which set the default value to `True`.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/modules/distributions/continuous.py b/torchrl/modules/distributions/continuous.py
--- a/torchrl/modules/distributions/continuous.py
+++ b/torchrl/modules/distributions/continuous.py
@@ -281,15 +281,15 @@ def log_prob(self, value, **kwargs):
 class TanhNormal(FasterTransformedDistribution):
     """Implements a TanhNormal distribution with location scaling.
 
-    Location scaling prevents the location to be "too far" from 0 when a TanhTransform is applied, which ultimately
-    leads to numerically unstable samples and poor gradient computation (e.g. gradient explosion).
-    In practice, the location is computed according to
+    Location scaling prevents the location to be "too far" from 0 when a
+    ``TanhTransform`` is applied, but ultimately
+    leads to numerically unstable samples and poor gradient computation
+    (e.g. gradient explosion).
+    In practice, with location scaling the location is computed according to
 
         .. math::
             loc = tanh(loc / upscale) * upscale.
 
-    This behaviour can be disabled by switching off the tanh_loc parameter (see below).
-
 
     Args:
         loc (torch.Tensor): normal distribution location parameter


Test patch:
Repository: pytorch/rl

Issue:
[BUG] Misleading error message in `advantage_module.set_keys()`
## Describe the bug

[`set_keys(value="new_value_key")`](https://github.com/pytorch/rl/blob/fa45472da29cc1d07b3ae57bdfe094207deeff13/torchrl/objectives/value/advantages.py#L218) expects the the `"new_value_key"` to be in the `out_keys` of `self.value_network`, but it raises an error complaining about the old key `self.tensor_keys.value`.

## Expected behavior

If it's the new key that should be in the value network's output then the error message should complain about the new key, not the old key.
Otherwise, if `set_keys()` changes the output key of the value network, then the error should be thrown when the old key is not in the value network's out keys.

## Screenshots

<img width="1300" alt="image" src="https://github.com/pytorch/rl/assets/37197319/90c3a390-05f6-4419-8d19-16d41bfff782">

## Reason and Possible fixes

The feature was introduced very recently in #1175.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [ ] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
@Blonck

Resolving patch:
diff --git a/torchrl/objectives/value/advantages.py b/torchrl/objectives/value/advantages.py
--- a/torchrl/objectives/value/advantages.py
+++ b/torchrl/objectives/value/advantages.py
@@ -218,7 +218,7 @@ def set_keys(self, **kwargs) -> None:
                 and (value not in self.value_network.out_keys)
             ):
                 raise KeyError(
-                    f"value key '{self.tensor_keys.value}' not found in value network out_keys"
+                    f"value key '{value}' not found in value network out_keys {self.value_network.out_keys}"
                 )
         if self._tensor_keys is None:
             conf = asdict(self.default_keys)


Test patch:
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -6028,6 +6028,10 @@ def test_set_keys(self, value, adv, kwargs):
         module.set_keys(value=value)
         assert module.tensor_keys.value == value
 
+        with pytest.raises(KeyError) as excinfo:
+            module.set_keys(unknown_key="unknown_value")
+            assert "unknown_value not found" in str(excinfo.value)
+
     @pytest.mark.parametrize(
         "adv,kwargs",
         [
Repository: pytorch/rl

Issue:
[BUG] Unsqueezing the error in ddpg
https://github.com/pytorch/rl/blob/3595c45eeb0c3e45c41bfff99ecbc973a7d1704f/torchrl/objectives/ddpg.py#L272

this line is not compatible with multi-agent, as it unsuquezees a dimension before the agents one.

Is there a workaround? is a `unsqueeze(-1)` sufficient?



Discussion:
Yeah should work!
Or maybe we can just remove it? What is the reason for that line?

Resolving patch:
diff --git a/torchrl/objectives/ddpg.py b/torchrl/objectives/ddpg.py
--- a/torchrl/objectives/ddpg.py
+++ b/torchrl/objectives/ddpg.py
@@ -11,10 +11,10 @@
 from typing import Tuple
 
 import torch
+
 from tensordict.nn import dispatch, make_functional, repopulate_module, TensorDictModule
 from tensordict.tensordict import TensorDict, TensorDictBase
 from tensordict.utils import NestedKey
-
 from torchrl.modules.tensordict_module.actors import ActorCriticWrapper
 from torchrl.objectives.common import LossModule
 from torchrl.objectives.utils import (
@@ -269,7 +269,6 @@ def forward(self, tensordict: TensorDictBase) -> TensorDict:
         """
         loss_value, td_error, pred_val, target_value = self._loss_value(tensordict)
         td_error = td_error.detach()
-        td_error = td_error.unsqueeze(tensordict.ndimension())
         if tensordict.device is not None:
             td_error = td_error.to(tensordict.device)
         tensordict.set(


Test patch:
Repository: pytorch/rl

Issue:
[Bug] Cannot instantiate 2 losses with different keys
In MARL different agent groups have different loss instances.

It is not currently possible to instantiate the same loss 2 times with different keys

```python
loss_1 = DQNLoss(value_network=nn.Linear(3, 3), action_space="one_hot")
loss_1.set_keys(reward="a")
assert loss_1.tensor_keys.reward == "a"
loss_2 = DQNLoss(value_network=nn.Linear(3, 3), action_space="one_hot")
loss_2.set_keys(reward="b")
assert loss_1.tensor_keys.reward == "a"  # Fails, it is now b
```


Discussion:

Resolving patch:
diff --git a/torchrl/objectives/common.py b/torchrl/objectives/common.py
--- a/torchrl/objectives/common.py
+++ b/torchrl/objectives/common.py
@@ -102,7 +102,6 @@ def tensor_keys(self) -> _AcceptedKeys:
 
     def __new__(cls, *args, **kwargs):
         cls.forward = set_exploration_type(ExplorationType.MODE)(cls.forward)
-        cls._tensor_keys = cls._AcceptedKeys()
         self = super().__new__(cls)
         return self
 
@@ -113,6 +112,7 @@ def __init__(self):
         self._value_estimator = None
         self._has_update_associated = {}
         self.value_type = self.default_value_estimator
+        self._tensor_keys = self._AcceptedKeys()
         self.register_forward_pre_hook(_updater_check_forward_prehook)
         # self.register_forward_pre_hook(_parameters_to_tensordict)
 


Test patch:
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -10128,6 +10128,15 @@ def test_single_call(self, has_target, value_key, single_call, detach_next=True)
         assert (value != value_).all()
 
 
+def test_instantiate_with_different_keys():
+    loss_1 = DQNLoss(value_network=nn.Linear(3, 3), action_space="one_hot")
+    loss_1.set_keys(reward="a")
+    assert loss_1.tensor_keys.reward == "a"
+    loss_2 = DQNLoss(value_network=nn.Linear(3, 3), action_space="one_hot")
+    loss_2.set_keys(reward="b")
+    assert loss_1.tensor_keys.reward == "a"
+
+
 class TestBuffer:
     # @pytest.mark.parametrize('dtype', (torch.double, torch.float, torch.half))
     # def test_param_cast(self, dtype):
Repository: pytorch/rl

Issue:
[BUG] Incorrect deprecation warning for ContinuousBox maximum
## Describe the bug

Calling `maximum` on `ContinuousBox` yields the following deprecation warning:
```
DeprecationWarning: <class 'torchrl.data.tensor_specs.ContinuousBox'>.maximum is going to be deprecated in favour of <class 'torchrl.data.tensor_specs.ContinuousBox'>.low
```

I assume this is incorrect and instead suggest `<class 'torchrl.data.tensor_specs.ContinuousBox'>.high`

## To Reproduce

Run the PPO tutorial in Colab, observe the warning in `ProbabilisticActor` definition cell

## Expected behavior

```
DeprecationWarning: <class 'torchrl.data.tensor_specs.ContinuousBox'>.maximum is going to be deprecated in favour of <class 'torchrl.data.tensor_specs.ContinuousBox'>.high
```

## System info

Colab default env, Pytorch 2.1, torchrl 0.2.0

## Reason and Possible fixes

```
@property
    def maximum(self):
        warnings.warn(
            f"{type(self)}.maximum is going to be deprecated in favour of {type(self)}.high",
            category=DeprecationWarning,
        )
        return self._high.to(self.device)

```

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
--- a/torchrl/data/tensor_specs.py
+++ b/torchrl/data/tensor_specs.py
@@ -384,7 +384,7 @@ def minimum(self):
     @property
     def maximum(self):
         warnings.warn(
-            f"{type(self)}.maximum is going to be deprecated in favour of {type(self)}.low",
+            f"{type(self)}.maximum is going to be deprecated in favour of {type(self)}.high",
             category=DeprecationWarning,
         )
         return self._high.to(self.device)


Test patch:
Repository: pytorch/rl

Issue:
[BUG] Incorrect default value for `normalize_advantage`.
## Describe the bug

The default value of `normalize_advantage` in `ClipPPOLoss` and `KLPENPPOLoss`  is set to `True` even though the documentation states that it should be `False`.
The base class `PPOLoss` uses the correct default value. 

## Expected behavior

Change the default value to `False` to match the documentation and the base class.

## Checklist

- [X] I have checked that there is no similar issue in the repo (**required**)
- [X] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/objectives/ppo.py b/torchrl/objectives/ppo.py
--- a/torchrl/objectives/ppo.py
+++ b/torchrl/objectives/ppo.py
@@ -774,7 +774,7 @@ def __init__(
         entropy_coef: float = 0.01,
         critic_coef: float = 1.0,
         loss_critic_type: str = "smooth_l1",
-        normalize_advantage: bool = True,
+        normalize_advantage: bool = False,
         gamma: float = None,
         separate_losses: bool = False,
         reduction: str = None,
@@ -1007,7 +1007,7 @@ def __init__(
         entropy_coef: float = 0.01,
         critic_coef: float = 1.0,
         loss_critic_type: str = "smooth_l1",
-        normalize_advantage: bool = True,
+        normalize_advantage: bool = False,
         gamma: float = None,
         separate_losses: bool = False,
         reduction: str = None,


Test patch:
Repository: pytorch/rl

Issue:
[BUG] Default Value Mismatch Between Documentation and Implementation of the ProbabilisticActor class
## Describe the bug
I'd like to express my gratitude for the swift and diligent maintenance of the torchrl library. Your efforts are greatly appreciated. 


By the way, it appears that there is a discrepancy between the documented default values and the actual default values implemented in the code.  I think this could potentially lead to confusion and incorrect usage by users who rely on the documentation for guidance.


- Documentation : https://github.com/pytorch/rl/blob/main/torchrl/modules/tensordict_module/actors.py#L161
- Implementation
   - https://github.com/pytorch/rl/blob/main/torchrl/modules/tensordict_module/probabilistic.py#L108
   - https://github.com/pytorch/rl/blob/main/torchrl/modules/tensordict_module/actors.py#L301

<br>

## To Reproduce

When running an Actor without setting "with set_exploration_type(ExplorationType.RANDOM):", 
the Actor does not perform exploration as expected. 

This issue becomes particularly problematic when implementing training code without using a data collector, 
leading to problems in  training and learning performance.


<br>

<br>

Thank you again for your continuous support and hard work on maintaining the torchrl library.



Discussion:

Resolving patch:
diff --git a/torchrl/modules/tensordict_module/actors.py b/torchrl/modules/tensordict_module/actors.py
--- a/torchrl/modules/tensordict_module/actors.py
+++ b/torchrl/modules/tensordict_module/actors.py
@@ -158,7 +158,7 @@ class ProbabilisticActor(SafeProbabilisticTensorDictSequential):
             the output value. Should be one of: 'InteractionType.MODE',
             'InteractionType.MEDIAN', 'InteractionType.MEAN' or
             'InteractionType.RANDOM' (in which case the value is sampled
-            randomly from the distribution). Defaults to is 'InteractionType.RANDOM'.
+            randomly from the distribution). Defaults to is 'InteractionType.MODE'.
 
             .. note:: When a sample is drawn, the :class:`ProbabilisticActor` instance will
               first look for the interaction mode dictated by the


Test patch:
Repository: pytorch/rl

Issue:
[BUG] `LazyTensorStorage` allocates storage on wrong device
https://github.com/pytorch/rl/blob/d59b8105cd7437cccb7c480f15ce6caf67f587c6/torchrl/data/replay_buffers/storages.py#L870-L873

As can be seen in the snippet, the storage first allocates everything on the data device and only later moves it to the proper device.

This is problematic as, if i request

```python
return TensorDictReplayBuffer(
            storage=LazyTensorStorage(
                100_000_000, # Say this is 30GB
                device="cpu",
            ),
            sampler=sampler,
            batch_size=100,
        )
```

it allocates the 30GB first on the data device (gpu) and then later moves it to cpu


Discussion:

Resolving patch:
diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py
--- a/torchrl/data/replay_buffers/storages.py
+++ b/torchrl/data/replay_buffers/storages.py
@@ -867,10 +867,10 @@ def max_size_along_dim0(data_shape):
             return (self.max_size, *data_shape)
 
         if is_tensor_collection(data):
-            out = data.expand(max_size_along_dim0(data.shape))
+            out = data.to(self.device)
+            out = out.expand(max_size_along_dim0(data.shape))
             out = out.clone()
             out = out.zero_()
-            out = out.to(self.device)
         else:
             # if Tensor, we just create a MemoryMappedTensor of the desired shape, device and dtype
             out = tree_map(


Test patch:
Repository: keras-team/keras-core

Issue:
`weighted_metrics` not showing up next to training progress bar
Tried porting this example (https://keras.io/examples/generative/gpt2_text_generation_with_kerasnlp/) to Keras Core. `accuracy` (passed as a `weighted_metric`) doesn't show up in the training bar (across backends).

Reproduced it here (with JAX): https://colab.research.google.com/drive/1XByfu9-vaFLnXdUjW6tLypMYy36C8LVx?usp=sharing
This is with legacy Keras: https://colab.research.google.com/drive/1fucRLYK_7ZDMRN5GURMgRIe80JknaFhj?usp=sharing


Discussion:
Thanks for the report. Can you repro with a regular Keras Core model (not KerasNLP)?
> Thanks for the report. Can you repro with a regular Keras Core model (not KerasNLP)?

Here you go: https://colab.research.google.com/drive/1NJsyb_IzR1Lnc2YD1U4AAV-mcutx6DM1?usp=sharing!
@mihirparadkar@gmail.com

Resolving patch:
diff --git a/keras_core/trainers/trainer.py b/keras_core/trainers/trainer.py
--- a/keras_core/trainers/trainer.py
+++ b/keras_core/trainers/trainer.py
@@ -44,7 +44,7 @@ def compile(
             )
         else:
             self._compile_loss = None
-        if metrics is not None:
+        if metrics is not None or weighted_metrics is not None:
             self._compile_metrics = CompileMetrics(
                 metrics, weighted_metrics, output_names=output_names
             )


Test patch:
diff --git a/keras_core/trainers/trainer_test.py b/keras_core/trainers/trainer_test.py
--- a/keras_core/trainers/trainer_test.py
+++ b/keras_core/trainers/trainer_test.py
@@ -109,6 +109,22 @@ def __init__(self, units):
         # And those weights are tracked at the model level
         self.assertEqual(len(model.metrics_variables), 6)
 
+        # Models with only weighted_metrics should have the same 3 metrics
+        model_weighted = ModelWithMetric(units=3)
+        model_weighted.compile(
+            optimizer=optimizers.SGD(),
+            loss=losses.MeanSquaredError(),
+            weighted_metrics=[metrics.MeanSquaredError()],
+        )
+        model_weighted.fit(
+            x,
+            y,
+            batch_size=2,
+            epochs=1,
+            sample_weight=np.ones(2),
+        )
+        self.assertEqual(len(model_weighted.metrics), 3)
+
     @parameterized.named_parameters(
         [
             ("eager", True, False, False),
Repository: keras-team/keras-core

Issue:
Using torch backend
Using PyTorch backend.
Epoch 1/3
32/32 ━━━━━━━━━━━━━━━━━━━━ 1s 4ms/step - mean_absolute_error: 0.4083 - loss: 0.2566
Epoch 2/3
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - mean_absolute_error: 0.3805 - loss: 0.2151
Epoch 3/3
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - mean_absolute_error: 0.3704 - loss: 0.2056
Epoch 1/5
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - loss: 0.2699 - mae: 0.4200
Epoch 2/5
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - loss: 0.2409 - mae: 0.3940
Epoch 3/5
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - loss: 0.2271 - mae: 0.3856
Epoch 4/5
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - loss: 0.2174 - mae: 0.3785
Epoch 5/5
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - loss: 0.2120 - mae: 0.3699
Epoch 1/3
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - mean_absolute_error: 0.7020 - loss: 0.3334
Epoch 2/3
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - mean_absolute_error: 0.4075 - loss: 0.1271
Epoch 3/3
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 4ms/step - mean_absolute_error: 0.3776 - loss: 0.1010
32/32 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - mean_absolute_error: 0.8608 - loss: 0.9672
Traceback (most recent call last):
  File "E:\custom_train_step_in_torch.py", line 483, in <module>
    gan.fit(dataloader, epochs=1)
  File "C:\Python_310\lib\site-packages\keras_core\src\utils\traceback_utils.py", line 123, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Python_310\lib\site-packages\keras_core\src\utils\module_utils.py", line 26, in initialize
    raise ImportError(
ImportError: This requires the tensorflow module. You can install it via `pip install tensorflow`


Discussion:
Sounds like you're trying to use a feature that requires TF. Can you try to run `keras.config.disable_traceback_filtering()` and run again so we can see the full traceback?
> Sounds like you're trying to use a feature that requires TF. Can you try to run `keras.config.disable_traceback_filtering()` and run again so we can see the full traceback?

Thanks. I've tried doing what you suggested, but the result is the same.
The original code is at https://github.com/keras-team/keras-io/blob/master/guides/keras_core/custom_train_step_in_torch.py.
`keras.config.disable_traceback_filtering()` is not supposed to fix the problem it is supposed to give you a longer stack trace. You can run it and post the full stack trace here.

Resolving patch:
diff --git a/keras_core/utils/module_utils.py b/keras_core/utils/module_utils.py
--- a/keras_core/utils/module_utils.py
+++ b/keras_core/utils/module_utils.py
@@ -14,9 +14,9 @@ def available(self):
         if self._available is None:
             try:
                 self.initialize()
+                self._available = True
             except ImportError:
                 self._available = False
-            self._available = True
         return self._available
 
     def initialize(self):


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
All tensors must be on devices[0]: 0
### 🐛 Describe the bug

For https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/resnet, when use `python -m torch.distributed.launch --nproc_per_node 2 --master_addr localhost --master_port 29500  run_resnet_cifar10_with_engine.py`, there is an error that `All tensors must be on devices[0]: 0`


### Environment

torch=1.8.1


Discussion:
Hi, did you build colossalai from source?
yes, using `nproc_per_node 1` is ok. but `nproc_per_node 2` will meet this error.
Ok, let me try to reproduce this error. May I know which GPU you are using and how many GPUs are available on your machine?
Tesla-V100 and 2GPUs are avaliable in my machine.
Got, let me try to reproduce this issue. I will get back to you soon!

Resolving patch:
diff --git a/colossalai/initialize.py b/colossalai/initialize.py
--- a/colossalai/initialize.py
+++ b/colossalai/initialize.py
@@ -348,12 +348,12 @@ def initialize(model: Union[nn.Module, List[nn.Module]],
                     "added even though not specified in the configuration",
                     ranks=[0])
         elif is_using_sequence():
-            model = DDP(model, process_group=gpc.get_group(ParallelMode.SEQUENCE_DP))
+            model = DDP(model, process_group=gpc.get_group(ParallelMode.SEQUENCE_DP), device_ids=[torch.cuda.current_device()])
             if verbose:
                 logger.info(
                     'Model is using torch.nn.parallel.DistributedDataParallel for Sequence Parallelism', ranks=[0])
         elif is_using_ddp() and not is_using_pp() and amp_mode != AMP_TYPE.NAIVE:
-            model = DDP(model, process_group=gpc.get_group(ParallelMode.DATA))
+            model = DDP(model, process_group=gpc.get_group(ParallelMode.DATA), device_ids=[torch.cuda.current_device()])
             if verbose:
                 logger.info(
                     'Model is using torch.nn.parallel.DistributedDataParallel for Data Parallelism', ranks=[0])


Test patch: