I got a ValueError when running "env = balloon_env.BalloonEnv()"

[eightnight2049]

File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/balloon_learning_environment/env/features.py", line 541, in _add_wind_features
assert 0.0 <= deviations[level] <= 1.00001, 'Uncertainty not in [0, 1].'
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

cuda：11.8
cudnn：0.86

jax：0.4.10
jaxlib：0.4.10
flax：0.4.2

[eightnight2049]

/root/miniconda3/envs/ble/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py:186: DeprecationWarning: ml_dtypes.float8_e4m3b11 is deprecated. Use ml_dtypes.float8_e4m3b11fnuz
return xc.batched_device_put(aval, sharding, xs, devices, committed) # type: ignore
/root/miniconda3/envs/ble/lib/python3.9/site-packages/jax/_src/array.py:515: DeprecationWarning: ml_dtypes.float8_e4m3b11 is deprecated. Use ml_dtypes.float8_e4m3b11fnuz
self._npy_value = self._single_device_array_to_np_array() # type: ignore

/root/miniconda3/envs/ble/lib/python3.9/site-packages/flax/core/scope.py:740: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
abs_value_flat = jax.tree_leaves(abs_value)
/root/miniconda3/envs/ble/lib/python3.9/site-packages/flax/core/scope.py:741: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
value_flat = jax.tree_leaves(value)

Traceback (most recent call last):
File "", line 1, in
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/gin/config.py", line 1605, in gin_wrapper
utils.augment_exception_message_and_reraise(e, err_str)
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/gin/utils.py", line 41, in augment_exception_message_and_reraise
raise proxy.with_traceback(exception.traceback) from None
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/gin/config.py", line 1582, in gin_wrapper
return fn(*new_args, **new_kwargs)
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/balloon_learning_environment/env/balloon_env.py", line 145, in init
self.arena = balloon_arena.BalloonArena(self._feature_constructor_factory,
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/balloon_learning_environment/env/balloon_arena.py", line 160, in init
self.reset(seed)
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/balloon_learning_environment/env/balloon_arena.py", line 183, in reset
return self.feature_constructor.get_features()
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/balloon_learning_environment/env/features.py", line 327, in get_features
self._add_wind_features(feature_vector)
File "/root/miniconda3/envs/ble/lib/python3.9/site-packages/balloon_learning_environment/env/features.py", line 541, in _add_wind_features
assert 0.0 <= deviations[level] <= 1.00001, 'Uncertainty not in [0, 1].'
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

[wangzexin21]

me too. Have you solved this problem yet?