Sub-routine segment masks are generated from the unnormalized distributions and combined with the predicted action logits (across slots) to obtain aggregated action sequence logits \(\mathbf {\hat{a}}\) .
For training, we use the PonderNet loss [1]} as our learning objective which comprises two terms.
The first being a weighted sum of reconstruction errors of the action sequence, where each term corresponds to the error obtained when using \(k\)  slots, and the weight is the probability that the trajectory contains \(k \in \lbrace 1, \dots , K\rbrace \)  sub-routines.
The second is a regularization cost given by a KL-divergence term between the posterior and prior distributions of the number of sub-routines in a trajectory.
