[BUG] Incorrect shape requirements for categorical spaces
Open
@vmoens is already working on this.
Since Feb 4, 2023.
bug
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 484
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
Describe the bug
Two bugs in one:
- The
Collectorexpects the action to be atorch.item()rather than atorch.Tensorwhencategorical_action_encoding=True. - When
categorical_action_encoding=Falsebut we pass incorrect shapes, the Collector will silently fail instead of crashing and emitting an error.
Bug 2 is quite insidious -- if your model outputs integer actions instead of one hot, the collector will always select the first action, regardless of what the model output. This can lead users to think that something is wrong with their algorithm.
To Reproduce
# Part 1
from torchrl.collectors import SyncDataCollector
import random
def policy(data):
# Action should be the same shape as reward, done, etc...
data["action"] = torch.randint(0, 2, data["done"].shape)
return data
env = GymEnv("CartPole-v0", categorical_action_encoding=True)
collector = SyncDataCollector(env, policy)
for data in collector:
pass
Traceback (most recent call last):
File "/Users/smorad/r4l_torch/examples/test_dqn.py", line 23, in <module>
for data in collector:
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/torchrl/collectors/collectors.py", line 535, in iterator
tensordict_out = self.rollout()
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/torchrl/collectors/collectors.py", line 658, in rollout
self._tensordict = self.env.step(self._tensordict)
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/torchrl/envs/common.py", line 380, in step
tensordict_out = self._step(tensordict)
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/torchrl/envs/gym_like.py", line 183, in _step
self._env.step(action_np)
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/gym/wrappers/time_limit.py", line 50, in step
observation, reward, terminated, truncated, info = self.env.step(action)
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/gym/wrappers/order_enforcing.py", line 37, in step
return self.env.step(action)
File "/Users/smorad/miniforge3/envs/torch_rl/lib/python3.9/site-packages/gym/envs/classic_control/cartpole.py", line 132, in step
assert self.action_space.contains(action), err_msg
AssertionError: array([1]) (<class 'numpy.ndarray'>) invalid
Calling .item() fixes this, but then the action is no longer a tensor.
def policy(data):
data["action"] = torch.randint(0, 2, data["done"].shape).item()
return data
# Part 2
from torchrl.collectors import SyncDataCollector
import random
def policy(data):
data["action"] = torch.randint(0, 2, data["done"].shape) # We are outputting shape=[1]
# Collector should crash because it expects shape=[2], but instead it will silently always pick the first action.
return data
env = GymEnv("CartPole-v0", categorical_action_encoding=False)
collector = SyncDataCollector(env, policy)
for data in collector:
pass # No error!! This should throw an error!
Expected behavior
- The collector should expect action to be a tensor of shape
(1,)like reward, done, etc. - The collector should error if the shape of a one hot vector is not the shape of the categorical space
System info
None 1.22.1 3.9.13 | packaged by conda-forge | (main, May 27 2022, 17:01:00)
[Clang 13.0.1 ] darwin
Additional context
Add any other context about the problem here.
Reason and Possible fixes
If you know or suspect the reason for this bug, paste the code lines and suggest modifications.
Checklist
- I have checked that there is no similar issue in the repo (required)
- I have read the documentation (required)
- I have provided a minimal working example to reproduce the bug (required)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.