Batch size in DQN custom model

RLlib
1

I did not fully understand how to deal with batch sizes when creating custom models. In my case, I am creating a custom model that extends DQNRecurrentModel and I want to define layers and nn blocks myself. Moreover, my observation is a quite complicated dictionary and I want to apply different layers to different parts of it, therefore I cannot rely on the tensors in obs_flat but I need to create them myself.

Very simple example:

class MyModel(DQNRecurrentModel):

    def __init__(self,
        # ... init arguments
    ):
    self.fully_connected = tf.keras.layers.Dense(128)

    def forward(self, input_dict, state, seq_lens):
        obs_piece = input_dict['obs']['my_key']
        obs_tensor = tf.constant(obs_piece)  # Does not have batch dimension
        embedding = self.fully_connected(obs_tensor)

Keras layers/models generally require the first input dimension to be the batch size. However, in the forward() of DQNRecurrentModel, I receive only one observation.

I can add the first batch dimension, i.e. obs_tensor = tf.expand_dims(obs_tensor), but I’m wondering if this is the correct way of handling this. My fear is that this will be inefficient or even incorrect when the model will be trained with a batch sampled from the replay buffer.

How should I consider the batch size in the custom model?
Thank you very much!

2

Actually, I realized that batch size is present, it just wasn’t the first dimension of the tensor shape due to how I stack the observation fields. I solved it by just setting the correct axis argument in tf.stack().

I apologize for the unnecessary question!

1 Like