Multi-head model functionality

  • High: It blocks me to complete my task.

Hi all, trying to reproduce some of the following paper:

The functionality I need currently is to be able to train a common network with multiple ‘heads’, with each head trained on a different task.

I can implement the multiple tasks using the TaskSettableEnv, but from that point i’m not sure. Two questions:

  1. is it possible to store experiences from different tasks separately in a reply buffer (say for the dqn replay buffer)
  2. if 1 is true, would I need to write a custom loss function to train each tasks data on a different head?