PaperlayerView exact revision on arXiv

Source-linked AI summary

Causally Correct Partial Models for Reinforcement Learning

Danilo J. Rezende, Ivo Danihelka, George Papamakarios, Nan Rosemary Ke, Ray Jiang, Theophane Weber, Karol Gregor, Hamza Merzic, Fabio Viola, Jane Wang, Jovana Mitrovic, Frederic Besse, Ioannis Antonoglou, Lars Buesing

arXiv:2002.02836v1cs.LGcs.AIstat.ML

TL;DR

High-dimensional observations can make full future modeling expensive, motivating partial models, but omitted observations used by a policy can confound those models. The paper proposes causally correct partial-model alternatives and reports theoretical and experimental validation for policy changes.

  • Problem

    Modeling all future observations can be costly or intractable, while partial models may omit observations used by the policy and fail under policy changes.

    Source: §1 · p. 1, §4 · p. 5

  • Method

    The paper proposes a family of partial models that uses causal structure in the agent’s state-to-action computation to support predictions robust to behavior-policy changes.

    Source: §1 · pp. 1–2, §4 · p. 5

  • Results

    The proposed modifications address correctness under behavior-policy changes and are validated theoretically and experimentally.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The modifications target policy-change robustness, not correctness or robustness to other types of environmental intervention.

    Source: §6 · p. 8

    A state-independent behavior policy can avoid the bias, but exploring interesting states with it requires enormous training data.

    Source: §2 · p. 2

Abstract

from arXiv · show

In reinforcement learning, we can learn a model of future observations and rewards, and use it to plan the agent's next actions. However, jointly modeling future observations can be computationally expensive or even intractable if the observations are high-dimensional (e.g. images). For this reason, previous works have considered partial models, which model only part of the observation. In this paper, we show that partial models can be causally incorrect: they are confounded by the observations they don't model, and can therefore lead to incorrect planning. To address this, we introduce a general family of partial models that are provably causally correct, yet remain fast because they do not need to fully model future observations.

1. Introduction

Model-based reinforcement learning uses learned models of future observations and rewards to support planning, but modeling all observations can be intractable. The paper explains how partial models can fail under policy changes and proposes causally informed alternatives.

1. Introduction

  • In model-based reinforcement learning, agents use sensory data to learn models predicting future observations from past actions and observations, which can support planning and other uses.Other listed uses include synthetic training data, credit assignment, representation learning, and exploration.Source: §1 · p. 1
  • Modeling high-dimensional observations can require generating images across multiple planning steps and may emphasize irrelevant visual details, motivating models that omit some observations.The paper calls models that neither condition on nor generate the full set of observed data partial models.Source: §1 · p. 1
  • Partial models can fail to predict correctly under a new policy when information used by the policy but omitted from the model acts as a confounder, potentially misleading planning.The paper frames this failure as a problem in causal reasoning.Source: §1 · p. 1
  • The paper clarifies the causal problem, asks what information makes partial models correct across policy changes, and proposes solutions studied in simple MDPs, MiniPacman, and 3D environments.The proposed family is examined empirically in illustrative environments.Source: §1 · p. 2

2. A simple example: FuzzyBear

In FuzzyBear, a reward model that omits the intermediate bear state can learn behavior-policy associations rather than rewards robust to policy changes. This can mislead planning, and even action-level causal correction does not capture the value of making decisions after observing the bear.

2. A simple example: FuzzyBear

  • The partial reward model conditions on the action sequence but omits the stochastic intermediate bear state, so its learned reward can depend on the behavior policy.If experience mostly hugs teddy bears and runs from grizzlies, the model associates visiting and hugging with reward near 1, although the correct average is 0.25.Source: §2 · p. 2
  • With the option to stay home for 0.6, planning from a model that predicts forest-and-hug reward near 1 can choose the suboptimal forest action, whose true maximum reward is 0.5.The model is not robust to changes in the behavior policy.Source: §2 · p. 2
  • Making the behavior policy independent of the bear state can avoid the bias, but exploring interesting states this way requires enormous training data.The paper presents this as a reason to seek a model robust to behavior-policy changes instead.Source: §2 · p. 2
  • Using action-level interventions gives causally correct planning but remains open-loop: the optimal policy observes whether the bear is teddy or grizzly before choosing an action.The intervention-based plan predicts 0.25 for visiting and hugging and 0 for visiting and running, while the optimal policy obtains 0.5.Source: §2 · p. 3

3. Background on causal reasoning

Causal reasoning distinguishes observational associations from predictions under interventions, which can differ when unobserved variables influence both predictor and target. Under suitable graph structure, observed variables can support intervention predictions, a connection the paper applies to reinforcement-learning models.

3. Background on causal reasoning

  • A predictor trained on observational pairs can estimate p(y|x) when the data-generation process is unchanged, but confounding can make this differ from p(y|do(x)).The do-operator sets x independently of its parents, changing the data-generation process.Source: §3 · p. 3
  • A backdoor variable can block confounder-to-covariate paths, and a frontdoor can block covariate-to-target paths, allowing intervention predictions from observed distributions under suitable conditions.The backdoor adjustment formula requires p(x|z) > 0.Source: §3 · pp. 3–4
  • A model is causally correct for an intervention set when its predictions remain accurate after interventions in that set.The passage defines this in terms of agreement between the model’s post-intervention distribution and the target distribution.Source: §3 · p. 4
  • In reinforcement learning, actions correspond to covariates, future observations or rewards to targets, and the policy state can confound predictions from a partial model.A node between the agent’s internal state and action in the policy’s computational graph can serve as a backdoor.Source: §3 · p. 4
  • Importance weighting can estimate an intervention-specific marginal, but requires refitting for each intervention and can have high variance when the intervention differs greatly from the behavior distribution.The alternative described fits p(y|z,x) and uses the adjustment formula.Source: §3 · p. 4

4. Learning causally correct models

Partial models can be fast, but omitting observations used by the policy makes those observations confounders when predicting future outcomes. The proposed Causal Partial Model uses an action backdoor as its partial view, enabling causal predictions under new policies without modeling all observations.

4. Learning causally correct models

  • A partial model generates intermediate partial views and actions to update its internal state, rather than conditioning each prediction on the full observation history.The autoregressive model instead feeds generated observations back into its state; overshooting models generate future observations without intermediate observation data.Source: §4 · pp. 4–5
  • Omitted past observations can confound action-conditioned predictions, making a learned partial model unreliable under policy changes and unsuitable for planning across policies.The omitted observations affect actions through the policy but are absent from the model's inputs.Source: §4 · p. 5
  • The backdoor zt can be chosen from the agent's computational graph, with agent state, policy probabilities, and intended action among the options.Agent state can be informative but costly to model when high-dimensional; intended actions offer simple modeling and low branching for discrete actions.Source: §4 · pp. 5–6
  • The partial view vt is set to the action backdoor zt, yielding a Causal Partial Model that supports simulation under a new policy using backdoor adjustment.The action is conditionally independent of the agent state given zt.Source: §4 · p. 5
  • The causal correction applies to partial models generally, and the experiments focus on a deterministic RNN backbone whose states condition observation generation.The observations are not fed back autoregressively in this model family.Source: §4 · p. 6

5. Experiments

Across the experiments, causal partial models support more reliable planning than non-causal partial models, which can be unrealistically optimistic or produce worse search. The paper examines this pattern through value iteration, learned-model search, and rollouts in a visually rich environment.

5.1. Value-iteration analysis on MDPs

  • The value-iteration analysis derives causal and non-causal models from MDP transitions and behavior policies, isolating planning quality from model accuracy.It evaluates optimal policies computed from those models across randomly generated policies.Source: §5 · p. 6, §5.1 · p. 7

5. Experiments

  • The experiments cover exact MDP-based value iteration, learned models evaluated through tree search, and model rollouts in a visually rich 3D environment.The learned-model experiments use expectimax or a MuZero-style MCTS variant.Source: §5 · pp. 6–7, §5.2 · p. 7

5.1. Value-iteration analysis on MDPs

  • In FuzzyBear value iteration, CPM's optimal value was at least the behavior policy's value and achievable in the real environment, while NCPM became unrealistically optimistic as behavior improved.For bad behavior policies, CPM could choose the opposite of the intended action; further improvement required searching in state s1 and retraining on improved-policy data.Source: §5.1 · p. 7

5.2. Experiments with Expectimax and MCTS

  • In MiniPacman, NCPMs achieved visibly smaller reward with expectimax or MCTS; on AvoidFuzzyBear, only the non-causal model failed to solve the task with expectimax.MuZero-style MCTS mitigated some negative effects of the NCPM, and CPM-based search found a much better policy than the pretrained behavior policy in AvoidFuzzyBear.Source: §5.2 · p. 7, Figure 5 · p. 8

5.3. Visually rich 3D Environment

  • In 3D T-Maze rollouts, NCPM was unrealistically optimistic, whereas CPM generated food at the rollout end with around 50% chance, matching the random placement of food on either side.The models used policy probabilities as the partial view in this environment.Source: §5.3 · p. 8

6. Conclusion

The paper characterizes partial-model failures through causal reasoning and proposes modifications validated theoretically and experimentally. These modifications address correctness under policy changes, but not other types of environmental intervention.

6. Conclusion

  • The paper proposes a modification to partial models that supports correct predictions under changes in behavior policy, validated theoretically and experimentally.The conclusion characterizes the modification as simple and effective.Source: §6 · p. 8
  • The proposed modifications do not address model correctness or robustness against other types of intervention in the environment.The authors identify these other interventions as future work.Source: §6 · p. 8
  • In the 3D T-Maze, CPM rollouts with generated z respect random food placement, whereas rollouts with forced z consistently generate food blocks under actions matching the ground-truth policy's partial view.The rollouts are conditioned on three previous frames without revealing food location.Source: Figure 6 · p. 8

Supplementary to Causally Correct Partial Models for Reinforcement Learning

The supplement describes causal adjustment formulas for interventions, including conditions under which observational data can identify interventional outcomes. Without additional structure, the joint distribution alone does not determine the effect of setting x.

A. Backdoor and frontdoor adjustment formulas

B. Derivation of causal correctness for the models in Figure 3

The derivation distinguishes partial models that reproduce interventional predictions when perfectly trained from a model that generally does not. The supplement then describes backdoor choices, training and simulation procedures, and properties and uses of causal partial models.

B. Derivation of causal correctness for the models in Figure 3

  • Models (c) and (e) are causally correct when perfectly trained, while model (d) is generally incorrect because its observational conditional need not equal the interventional conditional.For model (c), the learned observation factors match real-world conditionals; for model (e), the learned observation and z factors match their real-world counterparts.Source: B. Derivation of causal correctness… · p. 13

C. Additional choices of the backdoor zt

  • An empty backdoor is appropriate only when the behavior policy is state-independent, and then supports open-loop fixed-action plans rather than policies.An intermediate policy layer can serve as z, while combinations can add information such as downsampled, quantized input bits.Source: C. Additional choices of the backdoor zt · p. 14

D. Algorithms for training and simulating from the model

  • Training models targets such as rewards and returns and learns distributions for both y_t and z_t; simulation under a new policy samples actions using the partial view and updates model state.The training setup assumes a distributed actor-learner arrangement.Source: D. Algorithms for training and… · p. 14

E. Discussion of model properties

  • Causal partial models generate only a partial view, which can be small and easy to model; in stochastic environments they are described as fast, causally correct, and low-variance.Low simulation variance is attributed to not modeling or generating unimportant background distractors.Source: E. Discussion of model properties · p. 15
  • Causal partial models are invariant to the π(a_t|z_t) distribution, but can simulate only policies conditioned on the starting state, past actions, and partial views.Evaluating policies conditioned on different features requires collecting their trajectories and retraining the model.Source: E. Discussion of model properties · p. 15

F. MuZero properties on stochastic environments

On stochastic environments, non-causal models can make biased value estimates during planning. MuZero features and chance nodes can mitigate some negative effects, while clustered models offer a way to limit search branching.

F. MuZero properties on stochastic environments

  • A correct root value discourages opening real-world-suboptimal branches, and the learned policy prior can also suppress them when the branch action is not most probable.On AvoidFuzzyBear, MuZero has fewer problems when p(teddy) < 0.5.Source: F. MuZero properties on stochastic… · p. 15
  • Chance nodes increase branching, making search shallower and causing it to bootstrap more from the causally correct root value.Averaging values across search-tree nodes gives that root value only a small weight until many simulations have run.Source: F. MuZero properties on stochastic… · pp. 15–16

G. Tree search experiments

  • In MuZero-style search, chance-node values average child values under learned outcome probabilities, and unexpanded children bootstrap without extra network evaluations per simulation.The experiment uses intended actions as chance outcomes; the non-causal model ignores those outcomes, and the authors do not recommend pairing it with chance-node MCTS.Source: G. Tree search experiments · p. 16
  • The reported experiments used at least 8 independent random seeds, and plots show 95% confidence intervals.Hyperparameters were selected by trying multiples of 3 for learning rate, exploration probability, and policy cost weight.Source: G. Tree search experiments · p. 16, Table 3 · p. 16

G.1. Details of experiments on MDPs

  • Figure 7(b) compares expectimax models, including clustered-probability and clustered-observation variants that approximate modeling those quantities.The clustered variants are described in Appendix G.4.Source: §G.1 · p. 16

G.2. Details of experiments on MiniPacman

The MiniPacman experiments use a stochastic task with increasing ghost counts, while simple extensions improve model training and clustering can control chance-node branching. Clustering is not guaranteed to provide a sufficient backdoor when reconstruction remains imperfect.

G.2. Details of experiments on MiniPacman

  • MiniPacman starts with 2 ghosts and adds one after each specified block of 2 or 4 levels; drawing movement direction before each ghost makes the environment more Markovian.The experiments use the stochastic version released by Guez et al. (2019).Source: §G.2 · p. 16
  • When trained on pretrained-policy data, expectimax performed worse than on on-policy data and could not test its proposed action in the real environment.The pretrained policy used expectimax with search depth 1.Source: §G.2 · pp. 16–17

G.3. Simple extensions

  • For known partial-view distributions, training p(z_t|h_t) minimizes the exact KL from m(z_t|s_t); replay resamples z_t from a stored m(z_t|s_t) distribution and replays each trajectory 4 times.The resampled z_t comes from a posterior, with π(a_t|z_t) identified as the ε-exploration distribution.Source: §G.3 · p. 17

G.4. Models trained by clustering

  • Clustering the modeled layer yields a discrete component index z_t whose category count controls chance-node branching; the clustering loss uses only the needed mixture components.The index of the best component is used as the inferred z_t during training.Source: §G.4 · p. 17
  • A cluster index is not guaranteed to suffice as a backdoor when reconstruction loss is nonzero; observations with many distractors may require an unrealistically large number of components.Source: §G.4 · p. 17

H. Dyna experiments on AvoidFuzzyBear

In AvoidFuzzyBear, policy improvement from off-policy experience succeeds with a causal model, whereas the non-causal model earns less while predicting too much. The appendix describes a Dyna-style policy-gradient implementation and value derivations, alongside a FuzzyBear planning comparison.

H. Dyna experiments on AvoidFuzzyBear

  • The causal model achieves optimal reward and correctly predicts its value, while the non-causal model earns less than 0.6 but expects more than 0.6.Figure 7(a) shows the comparison after models are trained on sub-optimal behavior-policy data and the evaluation policy is trained inside the model.Source: H. Dyna experiments on AvoidFuzzyBear · p. 18, Figure 7 · p. 18

H.1. Dyna-style policy-gradient algorithm

  • The proposed Dyna-style policy-gradient implementation simulates experiences from a partial model and uses policy gradients to learn an improved policy as if those experiences were real.The authors present this implementation as a proof of principle and leave further exploration for future work.Source: §H.1 · p. 18

I.1. Optimal Value Derivations

  • The theoretical value derivations assume that MDP transition probabilities and the data-generating policy are accurately learned; action-conditioned reward predictions can introduce non-causal bias.The derivation notes that the state probability is affected by the action input when predicting reward.Source: §I.1 · p. 19

I.2. Planning with Non-Causal vs Causal Models on FuzzyBear

  • In FuzzyBear, the non-causal planner always hugs and expects reward 1, exceeding the MDP's maximum expected reward of 0.5; the causal planner samples intention and never overestimates expected reward.The causal planner's intention distinguishes the teddy-bear and grizzly-bear cases, supporting hug for the former and run for the latter.Source: §I.2 · p. 21, Figure 8 · p. 21

I.3. Learning with optimal policy and varying ε-exploration

Across exploration rates, causal-model evaluations remain stable while non-causal evaluations can be overly optimistic. On AvoidFuzzyBear, the causal model also chooses home under data from random policies; a conditional Dirichlet mixture is used when the backdoor is action probabilities.

I.3. Learning with optimal policy and varying ε-exploration

  • At ε = 0, V* NCPM(π*) is 1.0 against FuzzyBear's maximum expected reward of 0.5; as ε → 1, its evaluation approaches the causal evaluation based on a uniformly random policy.Source: §I.3 · p. 22
  • On AvoidFuzzyBear, the causal model's optimal value equals the maximum expected reward of 0.6 for any training policy, while causal values are constant across ε.The described Figure 9 analysis varies ε over (0, 1].Source: §I.3 · pp. 21–22, Figure 9 · p. 22

I.4. Sub-optimal behavior policies on AvoidFuzzyBear

  • With random-policy data on AvoidFuzzyBear, the causal model prefers staying home for any sampled intentions, while the non-causal model varies its overly optimistic evaluations and chooses the wrong action.Source: §I.4 · p. 22

J.1. Conditional Dirichlet Mixture

  • For a backdoor based on action probabilities, the model p(z_t|h_t) is a mixture with N_c Dirichlet components whose concentration parameters are produced by a ReLU MLP.The component parameters are α_k(h_t) = α softmax(f_k(h_t)); α is the total concentration parameter and N_a is the number of actions.Source: §J.1 · p. 22

J.2. Hyper Parameters and Training

The appendix describes training settings, analyzes rollout behavior as policy exploration changes, and compares model-rollout configurations.

J.2. Hyper Parameters and Training

  • The 3D experiment hyper-parameter ranges resemble Gregor et al. (2019), and each reported experiment used at least five random seeds.Table 5 lists the hyper-parameters used.Source: §J.2 · p. 22, Table 5 · p. 23
  • Training interleaved T-maze with a simple Food level to speed learning.On the Food level, the agent walked around and ate food blocks.Source: §J.2 · p. 22
  • Bar-plot error bars show standard error across five runs of 30 episodes each.Source: §J.2 · p. 22
  • Figure 10 scatters 500 randomly generated policies by environment evaluation and corresponding model-optimal evaluation V*.Source: Figure 10 · p. 23

J.3. Analysis of rollouts

  • Rollout analysis used five rollouts per episode after three observed frames, classifying food in the first 25 frames by RGB color matching with seven K-means clusters.Rollouts were generated after ceiling performance but before policy entropy reduced exploration too far.Source: §J.3 · p. 23
  • As policy entropy falls, CPM becomes overoptimistic like NCPM; non-zero ε-exploration can avoid this.Earlier in training, CPM generates both food and no-food outcomes.Source: Figure 11 · p. 24

K. Model Rollouts

  • The rollout figures show NCPM conditioned on forced actions and CPM conditioned on forced actions with generated or forced z.Table 6 marks environment-dependent variables in blue and model-state-computable variables in orange.Source: Table 6 · p. 25, Figure 12 · p. 26, Figure 13 · p. 27
Loading 2002.02836v1…