Training with Rollouts

[mattkiim]

component_losses.update(
    self.descent_loss(
        x, goal_mask, safe_mask, unsafe_mask, requires_grad=True
    )
)

Can someone explain how descent_loss works? I've been trying to train the model with the rollout loss, but I'm unsure how to set the size of the trajectory for which the rollout loss is computed.

From tracing the code, I suspect changing EpisodicDataModule's trajectories_per_episode and trajectory_length will set the number of rollouts and the trajectory rollout length. Is this correct?

Also, the rollout epochs are causing OOM (out of memory) issues. Is this expected?

Thank you!