TF 代理的奖励、step_type和折扣形式

问：

我正在尝试根据本教程中的代码使用 tf agent 训练代理。我目前正在定制py_environment供自己使用。除了与环境相关的代码外，其余代码与教程中的代码完全相同。

def compute_avg_return(environment, policy, num_episodes=10):

  total_return = 0.0
  for _ in range(num_episodes):

    time_step = environment.reset()
    episode_return = 0.0

    while not time_step.is_last():
      action_step = policy.action(time_step) # <----- error on this line
      time_step = environment.step(action_step.action)
      episode_return += time_step.reward
    total_return += episode_return

  avg_return = total_return / num_episodes
  return avg_return.numpy()[0]

compute_avg_return(eval_env, random_policy, num_eval_episodes)

第一次运行上述代码时出现以下错误：

ValueError: Received a mix of batched and unbatched Tensors, or Tensors are not compatible with Specs.  num_outer_dims: 1.
Saw tensor_shapes:
   TimeStep(
{'discount': TensorShape([1]),
 'observation': TensorShape([1, 50, 30]),
 'reward': TensorShape([1]),
 'step_type': TensorShape([1])})
And spec_shapes:
   TimeStep(
{'discount': TensorShape([]),
 'observation': TensorShape([1, 50, 30]),
 'reward': TensorShape([]),
 'step_type': TensorShape([])})

从错误日志中，我用于观察的形状是正确的，但仍然说它们不兼容。所以我认为问题出在“折扣”、“奖励”和“step_type”的形式上？

但是我该怎么办呢？我找不到任何可以告诉我如何定义/更改这些属性形状的东西。