//  [python]
def _compute_td_error_body(self, states, actions, next_states, rewards, dones):
        """
        Args:
            states:
            actions:
            next_states:
            rewards:
                Shape should be (batch_size, 1)
            dones:
                Shape should be (batch_size, 1)

        Returns:

        """
        batch_size = states.shape[0]
        actions = tf.cast(actions, dtype=tf.int32)

        # tf.print(tf.expand_dims(tf.range(batch_size), axis=1).shape,
        #          actions.shape,
        #          tf.concat(
        #              values=(tf.expand_dims(tf.range(batch_size), axis=1),
        #                      actions), axis=1).shape
        #          )

        with tf.device(self.device):
            rewards = tf.tile(
                tf.reshape(rewards, [-1, 1]),
                tf.constant([1, self.q_func.n_atoms]))  # [batch_size, n_atoms]
            not_done = 1.0 - tf.tile(
                tf.reshape(dones, [-1, 1]),
                tf.constant([1, self.q_func.n_atoms]))  # [batch_size, n_atoms]
            discounts = tf.cast(
                tf.reshape(self.discount, [-1, 1]), tf.float32)  # [1, 1]
            z = tf.reshape(
                self._z_list, [1, self.q_func.n_atoms])  # [1, n_atoms]
            z = rewards + not_done * discounts * z  # [batch_size, n_atoms]
            z = tf.clip_by_value(z, self._v_min, self._v_max)  # [batch_size, n_atoms]
            b = (z - self._v_min) / self._delta_z  # [batch_size, n_atoms]

            index_help = tf.expand_dims(
                tf.tile(
                    tf.reshape(tf.range(batch_size), [-1, 1]),
                    tf.constant([1, self.q_func.n_atoms])),
                -1)  # [batch_size, n_atoms, 1]
            u, l = tf.math.ceil(b), tf.math.floor(b)  # [batch_size, n_atoms]
            u_id = tf.concat(
                [index_help, tf.expand_dims(tf.cast(u, tf.int32), -1)],
                axis=2)  # [batch_size, n_atoms, n_action]
            l_id = tf.concat(
                [index_help, tf.expand_dims(tf.cast(l, tf.int32), -1)],
                axis=2)  # [batch_size, n_atoms, n_action]

            target_q_next_dist = self.q_func_target(
                next_states)  # [batch_size, n_action, n_atoms]

            if self._enable_double_dqn:
                # TODO: Check this implementation is correct
                target_q_next_dist = tf.gather_nd(
                    target_q_next_dist,
                    tf.concat(
                        [tf.reshape(tf.range(batch_size), [-1, 1]),
                         tf.reshape(actions, [-1, 1])],
                        axis=1))
            else:
                target_q_next_sum = tf.reduce_sum(
                    target_q_next_dist * self._z_list_broadcasted, axis=2)  # [batch_size, n_action]
                actions_by_target_q = tf.cast(
                    tf.argmax(target_q_next_sum, axis=1),
                    tf.int32)  # [batch_size,]
                target_q_next_dist = tf.gather_nd(
                    target_q_next_dist,
                    tf.concat(
                        [tf.reshape(tf.range(batch_size), [-1, 1]),
                         tf.reshape(actions_by_target_q, [-1, 1])],
                        axis=1))  # [batch_size, n_atoms]

            action_indices = tf.concat(
                values=(tf.expand_dims(tf.range(batch_size), axis=1),
                        actions), axis=1)
            current_q_dist = tf.gather_nd(
                self.q_func(states), action_indices)  # [batch_size, n_atoms]

            td_errors = tf.reduce_sum(
                target_q_next_dist * (u - b) * tf.math.log(tf.gather_nd(current_q_dist, l_id)) +
                target_q_next_dist * (b - l) * tf.math.log(tf.gather_nd(current_q_dist, u_id)),
                axis=1)

        return td_errors