ValueError: Output array has wrong dimensionality

[Yefanov]

Hello, I have encountered a problem running vector environments for TurnFaucet-v0 and ms1 tasks, even with sparse reward.

python -m mani_skill2.examples.demo_vec_env -e TurnFaucet-v0 --reward-mode sparse

2023-07-28 15:02:12,107 - mani_skill2 - INFO - Monkey patching gym.spaces.Box.__repr__
opts: []
env_kwargs: {}
[2023-07-28 15:02:13.568] [svulkan2] [error] GLFW error: X11: The DISPLAY environment variable is missing
[2023-07-28 15:02:13.568] [svulkan2] [warning] Continue without GLFW.
2023-07-28 15:02:15,828 - mani_skill2 - INFO - RenderServer is running at: localhost:43115
env VecEnv(<TimeLimit<TurnFaucetEnv<TurnFaucet-v0>>>)
Observation space Dict(agent:Dict(qpos:Box(-inf, inf, (9,), float32), qvel:Box(-inf, inf, (9,), float32), base_pose:Box(-inf, inf, (7,), float32)), extra:Dict(tcp_pose:Box(-inf, inf, (7,), float32), target_angle_diff:Box(-inf, inf, (1,), float32), target_joint_axis:Box(-inf, inf, (3,), float32), target_link_pos:Box(-inf, inf, (3,), float32)), camera_param:Dict(base_camera:Dict(extrinsic_cv:Box(-inf, inf, (4, 4), float32), cam2world_gl:Box(-inf, inf, (4, 4), float32), intrinsic_cv:Box(-inf, inf, (3, 3), float32)), hand_camera:Dict(extrinsic_cv:Box(-inf, inf, (4, 4), float32), cam2world_gl:Box(-inf, inf, (4, 4), float32), intrinsic_cv:Box(-inf, inf, (3, 3), float32))), image:Dict(base_camera:Dict(Color:Box(0.0, 1.0, (128, 128, 4), float32), Position:Box(-inf, inf, (128, 128, 4), float32)), hand_camera:Dict(Color:Box(0.0, 1.0, (128, 128, 4), float32), Position:Box(-inf, inf, (128, 128, 4), float32))))
Action space Box(-1.0, 1.0, (8,), float32)
Episode 0
Traceback (most recent call last):
  File "/opt/conda/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/examples/demo_vec_env.py", line 111, in <module>
    main()
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/examples/demo_vec_env.py", line 62, in main
    obs = env.reset()
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/vector/vec_env.py", line 279, in reset
    return self.reset_wait(indices=indices)
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/vector/vec_env.py", line 273, in reset_wait
    self._update_np_buffer(results, indices)
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/vector/vec_env.py", line 226, in _update_np_buffer
    return stack_obs(
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/vector/vec_env.py", line 425, in stack_obs
    ret[key] = stack_obs(_obs, space[key], buffer=_buffer)
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/vector/vec_env.py", line 425, in stack_obs
    ret[key] = stack_obs(_obs, space[key], buffer=_buffer)
  File "/opt/conda/lib/python3.9/site-packages/mani_skill2/vector/vec_env.py", line 428, in stack_obs
    return np.stack(obs, out=buffer)
  File "<__array_function__ internals>", line 180, in stack
  File "/opt/conda/lib/python3.9/site-packages/numpy/core/shape_base.py", line 433, in stack
    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
  File "<__array_function__ internals>", line 180, in concatenate
ValueError: Output array has wrong dimensionality

The same error replicates in the RL Colab after running this cell (modified only env_id and reward_mode):

from mani_skill2.vector import VecEnv, make as make_vec_env
num_envs = 2 # recommended value for Google Colab. If you have more cores and a more powerful GPU you can increase this
env: VecEnv = make_vec_env(
    "TurnFaucet-v0",
    num_envs,
    obs_mode=obs_mode,
    reward_mode="sparse",
    control_mode=control_mode,
)
env.seed(0)
obs = env.reset()
print("Base Camera RGB:", obs['image']['base_camera']['rgb'].shape)
print("Base Camera RGB device:", obs['image']['base_camera']['rgb'].device)
env.close()

2023-07-28 14:46:18,376 - mani_skill2 - INFO - RenderServer is running at: localhost:51711

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

[<ipython-input-18-6567b89c585b>](https://localhost:8080/#) in <cell line: 11>()
      9 )
     10 env.seed(0)
---> 11 obs = env.reset()
     12 print("Base Camera RGB:", obs['image']['base_camera']['rgb'].shape)
     13 print("Base Camera RGB device:", obs['image']['base_camera']['rgb'].device)

8 frames

[/usr/local/lib/python3.10/dist-packages/mani_skill2/vector/vec_env.py](https://localhost:8080/#) in reset(self, indices)
    277     def reset(self, indices=None):
    278         self.reset_async(indices=indices)
--> 279         return self.reset_wait(indices=indices)
    280 
    281     def step_async(self, actions: np.ndarray) -> None:

[/usr/local/lib/python3.10/dist-packages/mani_skill2/vector/vec_env.py](https://localhost:8080/#) in reset_wait(self, indices)
    271         self.waiting = False
    272         vec_obs = self._get_torch_observations()
--> 273         self._update_np_buffer(results, indices)
    274         vec_obs.update(deepcopy(self._obs_np_buffer))
    275         return vec_obs

[/usr/local/lib/python3.10/dist-packages/mani_skill2/vector/vec_env.py](https://localhost:8080/#) in _update_np_buffer(self, obs_list, indices)
    224         for i, obs in zip(indices, obs_list):
    225             self._last_obs_np[i] = obs
--> 226         return stack_obs(
    227             self._last_obs_np, self.non_image_obs_space, self._obs_np_buffer
    228         )

[/usr/local/lib/python3.10/dist-packages/mani_skill2/vector/vec_env.py](https://localhost:8080/#) in stack_obs(obs, space, buffer)
    423             _obs = [o[key] for o in obs]
    424             _buffer = None if buffer is None else buffer[key]
--> 425             ret[key] = stack_obs(_obs, space[key], buffer=_buffer)
    426         return ret
    427     elif isinstance(space, spaces.Box):

[/usr/local/lib/python3.10/dist-packages/mani_skill2/vector/vec_env.py](https://localhost:8080/#) in stack_obs(obs, space, buffer)
    423             _obs = [o[key] for o in obs]
    424             _buffer = None if buffer is None else buffer[key]
--> 425             ret[key] = stack_obs(_obs, space[key], buffer=_buffer)
    426         return ret
    427     elif isinstance(space, spaces.Box):

[/usr/local/lib/python3.10/dist-packages/mani_skill2/vector/vec_env.py](https://localhost:8080/#) in stack_obs(obs, space, buffer)
    426         return ret
    427     elif isinstance(space, spaces.Box):
--> 428         return np.stack(obs, out=buffer)
    429     else:
    430         raise NotImplementedError(type(space))

/usr/local/lib/python3.10/dist-packages/numpy/core/overrides.py in stack(*args, **kwargs)

[/usr/local/lib/python3.10/dist-packages/numpy/core/shape_base.py](https://localhost:8080/#) in stack(arrays, axis, out)
    431     sl = (slice(None),) * axis + (_nx.newaxis,)
    432     expanded_arrays = [arr[sl] for arr in arrays]
--> 433     return _nx.concatenate(expanded_arrays, axis=axis, out=out)
    434 
    435 

/usr/local/lib/python3.10/dist-packages/numpy/core/overrides.py in concatenate(*args, **kwargs)

ValueError: Output array has wrong dimensionality

[Yefanov]

Issue arises when the observation space contains a 0-dimensional parameter like target_angle_diff. In this case, this parameter is represented as either a 0-dimensional numpy array or just a float number. However, the stack_obs function expects a 1-dimensional array with shape (1,). A quick fix can be applied to the mani_skill2/vector/vec_env.py file by adding two lines:

def stack_obs(obs: Sequence, space: spaces.Space, buffer: Optional[np.ndarray] = None):
    if isinstance(space, spaces.Dict):
        ret = {}
        for key in space:
            _obs = [o[key] for o in obs]
            _buffer = None if buffer is None else buffer[key]
            ret[key] = stack_obs(_obs, space[key], buffer=_buffer)
        return ret
    elif isinstance(space, spaces.Box):
        if len(np.array(obs[0]).shape) == 0: # check shape
            obs = [np.array(o)[None] for o in obs] # add dimension
        return np.stack(obs, out=buffer)
    else:
        raise NotImplementedError(type(space))

[xuanlinli17]

@StoneT2000