'Reshape inputs from [time_length, batch_size, ...] to
[time_length * batch_size, ...].
This allows for computing the value estimate in one go.'

def reshape_batched_inputs(self, all_obs, all_actions, internal_policy_states, policy_logits):

    batch_size = tf.shape(all_obs[0])[1]
    time_length = tf.shape(all_obs[0])[0]
    reshaped_obs = []
    for (obs, (obs_dim, obs_type)) in zip(all_obs, self.env_spec.obs_dims_and_types):
        if self.env_spec.is_discrete(obs_type):
            reshaped_obs.append(tf.reshape(obs, [(time_length * batch_size)]))
        elif self.env_spec.is_box(obs_type):
            reshaped_obs.append(tf.reshape(obs, [(time_length * batch_size), obs_dim]))
    reshaped_prev_act = []
    reshaped_policy_logits = []
    for (i, (act_dim, act_type)) in enumerate(self.env_spec.act_dims_and_types):
        prev_act = all_actions[i]
        if self.env_spec.is_discrete(act_type):
            reshaped_prev_act.append(tf.reshape(prev_act, [(time_length * batch_size)]))
        elif self.env_spec.is_box(act_type):
            reshaped_prev_act.append(tf.reshape(prev_act, [(time_length * batch_size), act_dim]))
        reshaped_policy_logits.append(tf.reshape(policy_logits[i], [(time_length * batch_size), (-1)]))
    reshaped_internal_policy_states = tf.reshape(internal_policy_states, [(time_length * batch_size), self.internal_policy_dim])
    time_step = (float(self.input_time_step) * tf.expand_dims(tf.to_float((tf.range((time_length * batch_size)) / batch_size)), (-1)))
    return (time_step, reshaped_obs, reshaped_prev_act, reshaped_internal_policy_states, reshaped_policy_logits)
