pytorch / pytorch/rl

[BUG] Incorrect shape requirements for categorical spaces

Open
#894 2 comments 0 reactions 1 assignee View on GitHub

@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:

  1. The Collector expects the action to be a torch.item() rather than a torch.Tensor when categorical_action_encoding=True.
  2. When categorical_action_encoding=False but 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

  1. The collector should expect action to be a tensor of shape (1,) like reward, done, etc.
  2. 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.