OPD

1
for traj, y_teacher in offline_dataset: # x: (b, L)int. y_teacher (b, L)int.
2
    loss = KL(student(traj), teacher(traj)) # 一次性输出 len 个 dist. 需要一些 mask.

1
for prompt in prompts: # x: (b, prompt_len)
2
    traj = student.generate(prompt) # (b, L)
3
    loss = KL(student(cat(prompt, traj)), teacher(cat(prompt, traj))) # 需要一些 mask.