Trouble initialize Rainbow Agent

[ben-arnao]

RainbowAgent(sess=tensorflow.compat.v1.Session(),
                                  num_actions=2,
                                  stack_size=4,
                                  observation_shape=samples[0].shape,
                                  observation_dtype=tensorflow.float32)

I get error

ValueError: Dimension 0 in both shapes must be equal, but are 32 and 3232. Shapes are [32] and [3232]. for '{{node concat}} = ConcatV2[N=2, T=DT_INT64, Tidx=DT_INT32](strided_slice_4, strided_slice_3, concat/axis)' with input shapes: [32,1], [3232,1], [] and with computed input tensors: input[2] = <1>.
  In call to configurable 'DQNAgent' (<class 'dopamine.agents.dqn.dqn_agent.DQNAgent'>)
  In call to configurable 'RainbowAgent' (<class 'dopamine.agents.rainbow.rainbow_agent.RainbowAgent'>)

In this part of the code:

rainbowagent.py def _build_target_distribution(self)

    batch_indices = tf.range(tf.cast(batch_size, tf.int64))[:, None]
    # size of next_qt_argmax: batch_size x 2
    batch_indexed_next_qt_argmax = tf.concat(
        [batch_indices, next_qt_argmax], axis=1)

What am i doing wrong?

[nhewitt99]

What is the value of samples[0].shape? I'm able to run your example without issues and have tried multiple different tuples as the observation shape. Any glaring mistakes like wrong dimensions are caught in other parts of the code, so I'm not sure what the cause could be.