[RLlib] Better error handling when return shape from step() mismatch in utils._flatten_multidiscrete

[y-he2]

Description

Hi and thanks for the hardworks so far, and an amazing distributed RL library. Hope I came to the right place for RLlib improvement requests.

Im just gonna drop a quick one here as I encounter them on the road, and more details could be provided as per request. Currently the error message could be hard to use for tracing back the issue, when the returned dimension mismatches the defined observation_space element dimensions.

In this example:

• Cause of error: I defined a MultiDiscrete( np.one( 22 ) ) element in my observation space, then I accidentally returned a np.one( 41 ) vector in my gym.step() function.
• Current error handling:

...\gymnasium\spaces\utils.py", line 158, in _flatten_multidiscrete
    onehot[offsets[:-1] + x.flatten()] = 1
ValueError: operands could not be broadcast together with shapes (22,) (41,)

Assuming no previous knowledge this could leave the user in some confusion as to where to start to debug.

Debug code to trace the error I used in utils.py line 156:

try:
    onehot[offsets[:-1] + x.flatten()] = 1
except ValueError as err:
    print( "space:", space )
    print( "x:", x )
    print( "onehot:", onehot )
    print( "offsets[:-1]:", offsets[:-1] )
    print( "x.flatten():", x.flatten() )
    print( "offsets[:-1].shape:", offsets[:-1].shape )
    print( "x.flatten().shape:", x.flatten().shape )
    print( "space.nvec", space.nvec )
    print( "space.nvec.size", space.nvec.size )
    raise err

Useful info to display:

• Which element in which space
• Defined space shape
• Return vector shape from step()

Best regards, Ian

Use case

As the title and this could help the user to better understand/trace their error when observation element vector size mismatch in step() return and self.observation_space definition.

[cszhu]

This P2 issue has seen no activity in the past 2 years. It will be closed in 2 weeks as part of ongoing cleanup efforts.

Please comment and remove the pending-cleanup label if you believe this issue should remain open.

Thanks for contributing to Ray!