setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (256, 6) + inhomogeneous part.

[svha1977]

Hi

I get this error on a custom gymnasium wrapped environment with variable episode length.
I suspect that it might be near episode end where it cannot collect x samples.

I did strip away all the logic in the env - that is having chat gpt do it so there might be stray artifacts in there, but its gets to produce the same error as I my real life scenario

I did it in a Jupyter notebook, so it's not directly debuggable but close.
LightZeroBug.ipynb

@puyuan1996 As discussed on discord I here post a reproduction of the issue. I also did at printout of current numpy version and light zero version (0.2.0)

[puyuan1996]

Hello,

This error usually occurs because batch_rewards contains data with inconsistent shapes or types, making it impossible to convert into a numpy array.

I recommend wrapping the line batch_rewards = np.asarray(batch_rewards) with a try-except block to print the details of the data inside the list. This will help pinpoint the cause immediately.

You can modify the code as follows:

# Modification suggestion
try:
    batch_rewards = np.asarray(batch_rewards)
except Exception as e:
    print(f"[Debug Error] Convert batch_rewards to numpy failed: {e}")
    # Print list length
    print(f"List length: {len(batch_rewards)}")
    # Print details of the first and last elements to check for consistency
    if len(batch_rewards) > 0:
        print("First element details:")
        print(f"  Type: {type(batch_rewards[0])}")
        if hasattr(batch_rewards[0], 'shape'):
            print(f"  Shape: {batch_rewards[0].shape}")
        print(f"  Value: {batch_rewards[0]}")
        
        print("Last element details:")
        print(f"  Type: {type(batch_rewards[-1])}")
        if hasattr(batch_rewards[-1], 'shape'):
            print(f"  Shape: {batch_rewards[-1].shape}")
        print(f"  Value: {batch_rewards[-1]}")
        
        # Check if there are mixed types (e.g., some scalars, some arrays)
        types = set(type(x) for x in batch_rewards)
        print(f"All unique types in list: {types}")

    raise e  # Re-raise the exception to stop execution as expected

Running this will display the specific data characteristics causing the error in your console.

(Partially answered by Gemini 3.0 Pro)

[svha1977]

@puyuan1996 I did try to download latest nightly build, and install that in my notebook. I did your code addition, but I now get another error. Not sure if it's before or after the other one? But this seem to be another place inside the code not directly related to env / rewards at lease not directly or?

back (most recent call last)
Cell In[10], line 59
56 my_config.exp_name = "DummyGym-MuZero-Scalar"
57 max_env_step = int(5e5)
---> 59 train_muzero_with_gym_env(
60 [my_config, my_create_config],
61 seed=0,
62 max_env_step=max_env_step,
63 )

File ~/Downloads/LightZero-main/lzero/entry/train_muzero_with_gym_env.py:177, in train_muzero_with_gym_env(input_cfg, seed, model, model_path, max_train_iter, max_env_step)
174 break
176 # The core train steps for MCTS+RL algorithms.
--> 177 log_vars = learner.train(train_data, collector.envstep)
179 if cfg.policy.use_priority:
180 replay_buffer.update_priority(train_data, log_vars[0]['value_priority_orig'])

File ~/miniforge3/envs/rl/lib/python3.10/site-packages/ding/worker/learner/base_learner.py:167, in BaseLearner._time_wrapper..wrapper(*args, **kwargs)
165 def wrapper(*args, **kwargs) -> Any:
166 with self._wrapper_timer:
--> 167 ret = fn(*args, **kwargs)
168 self._log_buffer[var_type][var_name] = self._wrapper_timer.value
169 return ret

File ~/miniforge3/envs/rl/lib/python3.10/site-packages/ding/worker/learner/base_learner.py:227, in BaseLearner.train(self, data, envstep, policy_kwargs)
224 policy_kwargs = {}
226 # Forward
--> 227 log_vars = self._policy.forward(data, **policy_kwargs)
229 # Update replay buffer's priority info
230 if isinstance(log_vars, dict):

File ~/Downloads/LightZero-main/lzero/policy/muzero.py:399, in MuZeroPolicy._forward_learn(self, data)
394 action_batch = torch.from_numpy(action_batch).to(self._cfg.device).unsqueeze(-1).long()
395 data_list = [mask_batch, target_reward,
396 target_value, target_policy, weights
397 ]
398 [mask_batch, target_reward, target_value, target_policy,
--> 399 weights] = to_torch_float_tensor(data_list, self._cfg.device)
401 target_reward = target_reward.view(self._cfg.batch_size, -1)
402 target_value = target_value.view(self._cfg.batch_size, -1)

File ~/Downloads/LightZero-main/lzero/policy/utils.py:633, in to_torch_float_tensor(data_list, device)
631 return output_data_list
632 else:
--> 633 raise TypeError("The type of input must be np.ndarray or List[np.ndarray]")

TypeError: The type of input must be np.ndarray or List[np.ndarray]

Im not sure how to progress here?

Repo is this dummy env

import numpy as np
from gymnasium import Env, spaces
from typing import Dict, Optional, Any


class DummyLightZeroEnv(Env):


    metadata = {"render_modes": []}

    def __init__(
        self,
        obs_dim: int = 27,
        min_episode_len: int = 85,
        max_episode_len: int = 93,
        env_config: Optional[Dict] = None,
    ):
        super().__init__()
        self.env_config = env_config or {}
        self.obs_dim = int(obs_dim)
        self.min_episode_len = int(min_episode_len)
        self.max_episode_len = int(max_episode_len)

        # Observation & action spaces (same structure as your original env)
        self.observation_space = spaces.Dict(
            {
                "obs": spaces.Box(
                    low=-np.inf,
                    high=np.inf,
                    shape=(self.obs_dim,),
                    dtype=np.float32,
                ),
                "action_mask": spaces.Box(
                    low=0.0,
                    high=1.0,
                    shape=(27,),
                    dtype=np.float32,
                ),
                "to_play": spaces.Discrete(1),  # always 0
            }
        )
        self.action_space = spaces.Discrete(27)

        # Runtime state
        self.current_step: int = 0
        self.episode_len: int = 0

        # Prebuild constant obs/mask
        self._obs_vec = np.ones((self.obs_dim,), dtype=np.float32)
        self._action_mask = np.ones((27,), dtype=np.float32)

    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict] = None):
        super().reset(seed=seed)

        # Sample episode length between 85 and 93 (inclusive)
        self.episode_len = int(
            self.np_random.integers(
                low=self.min_episode_len,
                high=self.max_episode_len + 1,
            )
        )
        self.current_step = 0

        obs = {
            "obs": self._obs_vec.copy(),
            "action_mask": self._action_mask.copy(),
            "to_play": 0,
        }
        info: Dict[str, Any] = {}
        return obs, info

    def step(self, action: int):
        # Ensure action is within the space but otherwise ignore
        assert self.action_space.contains(action), f"Invalid action: {action}"

        reward = self.current_step
        self.current_step += 1
        terminated = self.current_step >= self.episode_len
        truncated = False
        

        obs = {
            "obs": self._obs_vec.copy(),
            "action_mask": self._action_mask.copy(),
            "to_play": 0,
        }

        info: Dict[str, Any] = {}
        return obs, float(reward), bool(terminated), bool(truncated), info

This wrapper

import gym
from gym import spaces
from gym.envs.registration import register
import numpy as np


class DummyGymAdapter(gym.Env):
    metadata = {"render_modes": []}

    def __init__(self):
        super().__init__()
        self._core = DummyLightZeroEnv()

        obs_len = self._core.observation_space["obs"].shape[0]

        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_len,), dtype=np.float32
        )
        self.action_space = self._core.action_space

    def reset(self):
        # gym old API: return obs only
        obs_dict, info = self._core.reset()
        obs = obs_dict["obs"]
        return obs

    def step(self, action):
        # core env: new API
        obs_dict, reward, terminated, truncated, info = self._core.step(action)
        
        # Hard guard: reward must be a scalar
        assert np.array(reward).shape == (), f"Env reward is not scalar: {reward} (shape {np.array(reward).shape})"
        
        done = bool(terminated or truncated)
        obs = obs_dict["obs"]
        print(reward)
        return obs, float(reward), done, info

    def render(self):
        return None

    def close(self):
        return None


def make_dummy_gym_env(**kwargs):
    return DummyGymAdapter()


register(
    id="DummyGym-v0",
    entry_point=make_dummy_gym_env,
)

This config

import copy
import numpy as np

from lzero.entry import train_muzero_with_gym_env
from zoo.classic_control.cartpole.config.cartpole_muzero_config import (
    main_config as cartpole_main_config,
    create_config as cartpole_create_config,
)

# ---------------------------------------------------------------------
# 1) Copy base configs
# ---------------------------------------------------------------------
my_config = copy.deepcopy(cartpole_main_config)
my_create_config = copy.deepcopy(cartpole_create_config)

# ---------------------------------------------------------------------
# 2) Infer obs_dim from the already defined make_dummy_gym_env()
# ---------------------------------------------------------------------
tmp_env = make_dummy_gym_env()          # uses DummyEnvAdapter defined in this notebook
obs_shape = tmp_env.observation_space.shape
obs_dim = int(np.prod(obs_shape))
tmp_env.close()

print("DummyGym-v0 obs_dim =", obs_dim)

# ---------------------------------------------------------------------
# 3) Point env sections to DummyGym-v0
# ---------------------------------------------------------------------
for cfg in (my_config, my_create_config):
    cfg.env.env_id = "DummyGym-v0"
    cfg.env.continuous = False
    cfg.env.manually_discretization = False

# parallel envs etc. only need to be set on the main config
my_config.env.collector_env_num = 4
my_config.env.evaluator_env_num = 2
my_config.env.n_evaluator_episode = 4
my_config.env.stop_value = 1e9

# ---------------------------------------------------------------------
# 4) Update *only* the main config's model obs/action sizes
# ---------------------------------------------------------------------
ACTION_SPACE_SIZE = 27  # your 3x3x3 joint action

m = my_config.policy.model
m.observation_shape = obs_dim     # scalar for MLP
m.action_space_size = ACTION_SPACE_SIZE

# Optional: disable reanalyze for simplicity (both configs have this field)
my_config.policy.reanalyze_ratio = 0.0
my_create_config.policy.reanalyze_ratio = 0.0

# ---------------------------------------------------------------------
# 5) Run MuZero training
# ---------------------------------------------------------------------
my_config.exp_name = "DummyGym-MuZero-Scalar"
max_env_step = int(5e5)

train_muzero_with_gym_env(
    [my_config, my_create_config],
    seed=0,
    max_env_step=max_env_step,
)

@puyuan1996 Any suggestions how to progress?

[puyuan1996]

It seems to be related to the environment's reward type. You are currently using a float, but LightZero typically expects an array (numpy array).

Please refer to the implementation here:
https://github.com/opendilab/LightZero/blob/main/zoo/atari/envs/atari_lightzero_env.py#L175

In your DummyGymAdapter's step function, try changing return obs, float(reward), done, info to return the reward as a numpy array. For example:

# Before
# return obs, float(reward), done, info

# Suggested change
return obs, np.array(reward).astype(np.float32), done, info  # or np.array([reward]).astype(np.float32)

Please try testing it again after this fix. Thanks for your patience

[svha1977]

@puyuan1996 Did try changing the adapters step function as suggested. It make no difference unfortunately. It's also not related to variance in episode steps. I print the dummy reward and with a simplified run with just 10 steps pr episode it return the same error.

Sample output

DummyGym-v0 obs_dim = 27
[01-04 22:59:11] INFO [RANK0]: DI-engine DRL Policy base_learner.py:360
MuZeroModelMLP(
(representation_network): RepresentationNetworkMLP(
(fc_representation): Sequential(
(0): Linear(in_features=27, out_features=128, bias=True)
(1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): GELU(approximate='tanh')
(3): Linear(in_features=128, out_features=128, bias=True)
)
(norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
)
(dynamics_network): DynamicsNetwork(
(fc_dynamics_1): Sequential(
(0): Linear(in_features=155, out_features=128, bias=True)
(1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=128, out_features=128, bias=True)
(4): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
(fc_dynamics_2): Sequential(
(0): Linear(in_features=128, out_features=128, bias=True)
(1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=128, out_features=128, bias=True)
(4): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
(fc_reward_head): Sequential(
(0): Linear(in_features=128, out_features=32, bias=True)
(1): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=32, out_features=601, bias=True)
)
)
(prediction_network): PredictionNetworkMLP(
(fc_prediction_common): Sequential(
(0): Linear(in_features=128, out_features=128, bias=True)
(1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=128, out_features=128, bias=True)
(4): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
(fc_value_head): Sequential(
(0): Linear(in_features=128, out_features=32, bias=True)
(1): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=32, out_features=601, bias=True)
)
(fc_policy_head): Sequential(
(0): Linear(in_features=128, out_features=32, bias=True)
(1): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=32, out_features=27, bias=True)
)
)
(projection): Sequential(
(0): Linear(in_features=128, out_features=1024, bias=True)
(1): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=1024, out_features=1024, bias=True)
(4): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
(6): Linear(in_features=1024, out_features=1024, bias=True)
(7): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(prediction_head): Sequential(
(0): Linear(in_features=1024, out_features=512, bias=True)
(1): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Linear(in_features=512, out_features=1024, bias=True)
)
)
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
[01-04 22:59:12] WARNING The data in replay_buffer is not sufficient to sample a mini-batch: batch_size: 256, current buffer statistics train_muzero_with_gym_env.py:168
is: num_of_all_collected_episodes: 8, num of game segments: 8, number of transitions: 80 continue to collect
now ....
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
[01-04 22:59:12] WARNING The data in replay_buffer is not sufficient to sample a mini-batch: batch_size: 256, current buffer statistics train_muzero_with_gym_env.py:168
is: num_of_all_collected_episodes: 16, num of game segments: 16, number of transitions: 160 continue to collect
now ....
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
[01-04 22:59:13] WARNING The data in replay_buffer is not sufficient to sample a mini-batch: batch_size: 256, current buffer statistics train_muzero_with_gym_env.py:168
is: num_of_all_collected_episodes: 24, num of game segments: 24, number of transitions: 240 continue to collect
now ....
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
9.0

TypeError Traceback (most recent call last)
Cell In[12], line 59
56 my_config.exp_name = "DummyGym-MuZero-Scalar"
57 max_env_step = int(5e5)
---> 59 train_muzero_with_gym_env(
60 [my_config, my_create_config],
61 seed=0,
62 max_env_step=max_env_step,
63 )

File ~/Downloads/LightZero-main/lzero/entry/train_muzero_with_gym_env.py:177, in train_muzero_with_gym_env(input_cfg, seed, model, model_path, max_train_iter, max_env_step)
174 break
176 # The core train steps for MCTS+RL algorithms.
--> 177 log_vars = learner.train(train_data, collector.envstep)
179 if cfg.policy.use_priority:
180 replay_buffer.update_priority(train_data, log_vars[0]['value_priority_orig'])

File ~/miniforge3/envs/rl/lib/python3.10/site-packages/ding/worker/learner/base_learner.py:167, in BaseLearner._time_wrapper..wrapper(*args, **kwargs)
165 def wrapper(*args, **kwargs) -> Any:
166 with self._wrapper_timer:
--> 167 ret = fn(*args, **kwargs)
168 self._log_buffer[var_type][var_name] = self._wrapper_timer.value
169 return ret

File ~/miniforge3/envs/rl/lib/python3.10/site-packages/ding/worker/learner/base_learner.py:227, in BaseLearner.train(self, data, envstep, policy_kwargs)
224 policy_kwargs = {}
226 # Forward
--> 227 log_vars = self._policy.forward(data, **policy_kwargs)
229 # Update replay buffer's priority info
230 if isinstance(log_vars, dict):

File ~/Downloads/LightZero-main/lzero/policy/muzero.py:399, in MuZeroPolicy._forward_learn(self, data)
394 action_batch = torch.from_numpy(action_batch).to(self._cfg.device).unsqueeze(-1).long()
395 data_list = [mask_batch, target_reward,
396 target_value, target_policy, weights
397 ]
398 [mask_batch, target_reward, target_value, target_policy,
--> 399 weights] = to_torch_float_tensor(data_list, self._cfg.device)
401 target_reward = target_reward.view(self._cfg.batch_size, -1)
402 target_value = target_value.view(self._cfg.batch_size, -1)

File ~/Downloads/LightZero-main/lzero/policy/utils.py:633, in to_torch_float_tensor(data_list, device)
631 return output_data_list
632 else:
--> 633 raise TypeError("The type of input must be np.ndarray or List[np.ndarray]")

TypeError: The type of input must be np.ndarray or List[np.ndarray]

[SteamMachinist]

I'm experiencing the same issue, running code in kaggle though. Can share the code if it will be useful.

[svha1977]

I seem to get it to work in a dummy sense. My findings is

1. I made a pure LightZero env, not a gym wrapper. Essentially like in my gym dummy env just returning simple values and no real logic
2. Did try with latest main branch - i.e. or the one I downloaded a month or so ago
3. Same error as above
4. Did try with latest release by pip install that 0.2.0 version and that for now did work.
5. Will see if something pop up when doing real stuff.

So it seem there is some bug introduced after the 0.2.0 release currently in the main branch. And somehow the original bug here went away by using the pure LightZero env, so maybe that first error is related to the wrapper integration

[SteamMachinist]

I think the issue is fixed with new commits from mantainer(s), I managed to start training without any problems after fixing few small bugs. You can use my fork or wait untill it will hopefully be merged: #463