Objective-Aligned Amortized Inference for Offline Bayes-Adaptive MDP Model Learning

Learning latent environment representations from offline datasets is an important challenge for Bayes-adaptive Markov decision process model learning. A key difficulty is that offline datasets collected under heterogeneous protocols can differ in state--action visitation even when the underlying dynamics are identical. Visitation-decoupled objectives address this by evaluating transition predictive fit under a protocol-independent reference distribution via a change of measure. We show, however, that the resulting objective-induced invariances need not be inherited by amortized inference. An unconstrained encoder can depend on protocol-induced visitation variations that are invisible to the objective, a structural failure mode we term \emph{amortization mismatch}. To address this issue, we formalize an \emph{objective-aligned} design principle as an objective--inference compatibility condition: the encoder should factor through the objective-induced dataset signature. We instantiate this principle in multiple set-encoder families and, through targeted experiments, provide evidence that objective-aligned encoders improve posterior consistency and can mitigate downstream planning discrepancy across protocols.