fix

markub3327 · Nov 17, 2023 · c8ce6f0 · c8ce6f0
1 parent c903de9
commit c8ce6f0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rl_toolkit/agents/dueling_dqn/server.py b/rl_toolkit/agents/dueling_dqn/server.py
@@ -143,7 +143,7 @@ def __init__(
                             [*self._env.action_space.shape],
                             self._env.action_space.dtype,
                         ),
-                        "ext_reward": tf.TensorSpec([frame_stack], tf.float64),
+                        "ext_reward": tf.TensorSpec([frame_stack, 1], tf.float64),
                         "next_observation": tf.TensorSpec(
                             [*self._env.observation_space.shape],
                             self._env.observation_space.dtype,