[RLlib] Repeated space: The two structures don't have the same number of elements

[Misterion777]

What happened + What you expected to happen

Hi, I’m encountering following error:

ValueError: The two structures don't have the same nested structure.
...
More specifically: The two structures don't have the same number of elements.
...

during compute_single_action() (see script). To my understanding it is happening during mapping incoming observation to one sampled from the observation space. My observation space has a Repeated subspace.
As I can see, Sampling from Repeated space produces a list with random length.
Mapping uses tree.assert_same_structure method which raises error when:

If the two structures do not have the same number of elements or if the two structures are not nested in the same way.

It seems, that randomizing number of elements when sampling from Repeated is incorrect behaviour in this case as it will definitely not have the same number of elements with real observation, therefore this method will most likely raise an error, although I can be wrong.

Versions / Dependencies

Python: 3.8.12
Ray: 1.13.0
Torch: 1.11.0
OS: MacOS

Reproduction script

import numpy as np
import gym
import ray
from ray.rllib.agents import dqn       
from ray.tune.registry import register_env
from gym.spaces import Discrete, Box, Dict
from ray.rllib.utils.spaces.repeated import Repeated

MAX_ELEMENTS = 50
X_SHAPE = (5,128)
Y_SHAPE = (2,)
Z_SHAPE = (1,)
CUSTOM_SPACE = Dict({
            "x": Repeated(Box(-1,1,shape=X_SHAPE,dtype=np.float32), max_len=MAX_ELEMENTS),
            "y": Repeated(Box(0,MAX_ELEMENTS,shape=Y_SHAPE,dtype=np.int64), max_len=2 * MAX_ELEMENTS),
            "z": Repeated(Discrete(3), max_len=MAX_ELEMENTS),
    }
)
CUSTOM_SAMPLES = [CUSTOM_SPACE.sample() for _ in range(10)]
class NestedDictEnv(gym.Env):
    def __init__(self):
        self.action_space = Discrete(2)
        self.observation_space = CUSTOM_SPACE        
        self.steps = 0

    def reset(self):
        self.steps = 0
        return CUSTOM_SAMPLES[0]

    def step(self, action):
        self.steps += 1
        return CUSTOM_SAMPLES[self.steps], 1, self.steps >= 5, {}


ray.init()
register_env("custom_nested", lambda _: NestedDictEnv())
trainer = dqn.DQNTrainer(config={
    "env": "custom_nested",        
    "framework": "torch"   
})        

example_obs = {
    "x": list(np.zeros((10,) + X_SHAPE,dtype=np.float32)),
    "y": list(np.zeros((20,) + Y_SHAPE,dtype=np.int64)),
    "z": list(np.zeros((10,) + Z_SHAPE,dtype=np.int64))
}
trainer.compute_single_action(example_obs)

Issue Severity

High: It blocks me from completing my task.

[kouroshHakha]

Hi @Misterion777 the reason for this issue is that in your script the example_obs does not match the structure specified by the space of your environment. To test that this is really the case you can use env.observation_space.sample() or env.reset() to get a proper obs struct and then your script would work.

env = NestedDictEnv()
obs = env.reset()
act = trainer.compute_single_action(obs)