While the methodology discussed so far can produce effective policies, language models also represent task dynamics, as discussed in Section . We can leverage this fact to further improve the performance of our fine-tuned models by performing model-based planning at test-time, using both the policy and dynamics components in concert to further maximize task reward. A full dialogue trajectory can then be formed by concatenating this sampled future trajectory \(\tau _{\ge t}\)  with the current state of the dialogue \(\tau _{<t}\)  i.e., \(\tau =\lbrace \tau _{<t}, \tau _{\ge t}\rbrace \) . We perform the model-based planning by sampling \(k\)  such future trajectories from the final fine-tuned model, and ranking them according to an estimated reward function \(\hat{R}(\tau , c_o)\)  (see Appendix REF ). Then, we improve upon the policy \(\pi \)  from which we took the samples by taking the action (i.e., the next utterance) \(a_t\)  which receives the highest estimated reward among the sampled trajectories. This rollout sampling procedure is identical to the one used by [1]}.
