Why only one timestep for the imitation loss?

[telejesus2]

Hi! I have noticed that during training the control network has to predict only one action per episode (one from a query demo and one from a support demo so two in total).

I am curious what is the goal of this? Is it a way to balance the training of the embedding and control networks? Did you try other strategies (e.g. predicting all actions of all demos in the episode, or something in between)?

Thanks!