We want our model to dynamically predict the keyframe placement τ. However, learning a distribution over this discrete variable is challenging due to the need of discrete sampling. To learn the keyframe placement efficiently and in a differentiable manner, we propose a continuous relaxation of the objective.

In our soft relaxation of the objective, instead of sampling from τ and computing the loss on the corresponding target, we use the expected target in the loss. As shown to the left, the target for each keyframe is simply the weighted average of ground truth frames under the corresponding distribution τ. The keyframe prediction loss is then the reconstruction loss between the prediction and this target.

We can now use our keyframe model to perform hierarchical planning. To construct a plan, we search for a sequence of keyframes that starts at the current observed image and reaches the goal. We use the cross-entropy method in the latent space to optimize the sequence.

Once we constructed the plan expressed in terms of keyframes, we treat them as subgoals, and pass them to a low-level controller, e.g. based on model predictive control. We use the controller to produce action trajectories that reach each keyframe sequentially, until the final goal is reached.