python - 如何在 step 方法中为 TF-Agents 环境提供一组动作？

Question

我想对我的环境进行建模，使每个动作都由 3 个可能的子动作组成。

我将我的定义_action_spec为tf_agents.environments.py_environment.PyEnvironment：

self._action_spec = tf_agents.specs.BoundedArraySpec(
            shape=(3,), dtype=np.int32, name="action", minimum=[0, 0, 0], maximum=[10, 11, 12])

我的step方法失败了，我正在尝试：

env = NetworkEnv(discount=0.9)
tf_env = TFPyEnvironment(env)
print(tf_env.reset())
action = tf.constant([3, 3, 3], dtype=tf.int32, shape=(3,), name='action')
print(tf_env.step(action))
tf_env.close()

但它给ValueError: cannot select an axis to squeeze out which has size not equal to one

我应该如何为step方法提供操作？