Repository: pytorch/rl

Issue:
[BUG] I had to patch these 2 methods in order to run my script
## Describe the bug

(1)
`DoubleToFloat` transform cannot transform to `float` if env to be transformed is already a `float` (e.g. it is a `float` and the input dtype was declared as `double`). 

I mean, all it has to do is call `.to(torch.float)` or something, which can be called regardless of the `dtype` of the input tensor, so why not?

(2)
The `torchrl.objectives.utils._cache_values` function is also bugged I think. If there is no `_cache` attribute in `self.__dict__`, it raises an error, but the function seems to be implemented in a way that it can work anyway. In fact, when I manually do `self.__dict__["_cache"] = {}`, all is good.

This happens with the `torchrl.objectives.ClipPPOLoss`.

## To Reproduce

I do not know how to reproduce, although my project is fully available at: https://gitlab.com/svnv-svsv-jm/chess-rl/-/tree/develop?ref_type=heads

For more context, I am trying to combine `torchrl` with `pytorch-lightning`.

## Reason and Possible fixes

For the cache issue with `torchrl.objectives.ClipPPOLoss`, I patched this:

```python
# in torchrl.objectives.utils
def _cache_values(fun):
    """Caches the tensordict returned by a property."""
    name = fun.__name__

    def new_fun(self, netname=None):
        __dict__ = self.__dict__
       # MY PATCH START
        try:
            _cache = __dict__["_cache"]
        except KeyError as ex:
            __dict__["_cache"] = {}
            _cache = __dict__["_cache"]
        # MY PATCH END
        attr_name = name
        if netname is not None:
            attr_name += "_" + netname
        if attr_name in _cache:
            out = _cache[attr_name]
            return out
        if netname is not None:
            out = fun(self, netname)
        else:
            out = fun(self)
        # TODO: decide what to do with locked tds in functional calls
        # if is_tensor_collection(out):
        #     out.lock_()
        _cache[attr_name] = out
        return out

    return new_fun
```

For the `torch.envs.DoubleToFloat`, I patched `DTypeCastTransform`:

```python
# class DTypeCastTransform
def transform_observation_spec(self, observation_spec):
        full_observation_spec = observation_spec
        for observation_key, observation_spec in list(
            full_observation_spec.items(True, True)
        ):
            # find out_key that match the in_key
            for in_key, out_key in zip(self.in_keys, self.out_keys):
                if observation_key == in_key:
                    if observation_spec.dtype != self.dtype_in:
                        # PATCH START
                        self.dtype_in = observation_spec.dtype
                        # raise TypeError(
                        #     f"observation_spec.dtype is not {self.dtype_in}"
                        # )
                        # PATCH END
                    full_observation_spec[out_key] = self._transform_spec(
                        observation_spec
                    )
        return full_observation_spec
```

After these patches, all runs smoothly again.

Sorry for not providing a minimal example, but I really wouldn't know how to create one from my code. These are the tests where I encountered the errors: https://gitlab.com/svnv-svsv-jm/chess-rl/-/tree/develop/tests/shark/models/ppo?ref_type=heads



Discussion:
Thanks for reporting!

For the first point, I'd like to understand if your specs are accurate or not? It seems to me that you're asking that we can pass keys which specs are already float to it right?
Why not just doing `TransformedEnv(env, DoubleToFloat())` and let `DoubleToFloat` detect all the double and transform them?

For the second point, I think we can remove the caching mechanism in many places. Your patch could be replaced by a `setdefault` no? I can do that easily :)
Hi! Thanks for the reply.

Yes that's what I am asking for :) For more info, I am not using a custom env there, I was testing around with "InvertedDoublePendulum-v4".

So as regards the `DoubleToFloat` issue, I think just calling `DoubleToFloat()` with args may suffice. I wasn't aware that it would've made a difference if I called it with some input keys or without.

For the cache things, yes exactly I think so, too. That function can already work with `__dict__['_cache'] == {}`, so it'd make sense to make sure that the `__dict__` dictionary attribute has a default value of `{}` for the key `"_cache"`.
PS: if I am not wrong, my training loop works fine if run it as plain pytorch code, the cache-related bug arose when calling everything via the pytorch lightning trainer. Not sure why! Maybe this info is useful, otherwise ignore it.
There's probably a deepcopy of some sort somewhere!
I'd be interested to see where things break with lightning, there is certainly a lot that can be done on torchrl side for a better integration

Resolving patch:
diff --git a/torchrl/objectives/utils.py b/torchrl/objectives/utils.py
--- a/torchrl/objectives/utils.py
+++ b/torchrl/objectives/utils.py
@@ -459,7 +459,7 @@ def _cache_values(fun):
 
     def new_fun(self, netname=None):
         __dict__ = self.__dict__
-        _cache = __dict__["_cache"]
+        _cache = __dict__.setdefault("_cache", {})
         attr_name = name
         if netname is not None:
             attr_name += "_" + netname


Test patch: