[Paper Review] Counterfactual Control for Free from Generative Models
This paper proposes a method to derive control policies on the fly from a generative model of action-state trajectories, using gradient descent in the model's latent space to optimize for any reward function without retraining. The key contribution is enabling dynamic goal adaptation and counterfactual planning via a single pre-trained generative model, demonstrated by successfully balancing and tracking a moving target in CartPole without direct training on the task.
We introduce a method by which a generative model learning the joint distribution between actions and future states can be used to automatically infer a control scheme for any desired reward function, which may be altered on the fly without retraining the model. In this method, the problem of action selection is reduced to one of gradient descent on the latent space of the generative model, with the model itself providing the means of evaluating outcomes and finding the gradient, much like how the reward network in Deep Q-Networks (DQN) provides gradient information for the action generator. Unlike DQN or Actor-Critic, which are conditional models for a specific reward, using a generative model of the full joint distribution permits the reward to be changed on the fly. In addition, the generated futures can be inspected to gain insight in to what the network 'thinks' will happen, and to what went wrong when the outcomes deviate from prediction.
Motivation & Objective
- To develop a control framework that decouples policy learning from reward specification, enabling dynamic goal adaptation.
- To enable the agent to plan and select actions based on counterfactual futures imagined by a generative model of action-state trajectories.
- To demonstrate that a single generative model can support control for multiple, unseen tasks without retraining.
- To explore whether joint generative modeling of actions and future states enables more robust and generalizable control than task-specific models.
Proposed method
- The method uses a recurrent variational autoencoder (RVAE) to learn the joint distribution $p(future, action | state)$ from interaction data.
- Control policy is derived by optimizing over the latent space of the RVAE using gradient descent to maximize a desired reward function on generated counterfactual futures.
- The reward function is evaluated directly on generated future trajectories, with gradients backpropagated through the generative model to update the action sequence.
- The approach leverages the full joint distribution to ensure action-state sequences are internally consistent and plausible.
- The model is trained once on environment interactions and can then support any reward function via latent space search.
- The method avoids training a separate reward network or policy head, relying instead on the generative model’s internal dynamics for gradient estimation.
Experimental results
Research questions
- RQ1Can a single generative model of action-state trajectories be used to induce control policies for arbitrary reward functions without retraining?
- RQ2Can the model generalize to novel tasks not seen during training, such as tracking a moving target?
- RQ3How sensitive is the method to hyperparameters, particularly the latent space dimension?
- RQ4Can the model generate internally consistent counterfactual futures that support effective planning and control?
Key findings
- The model successfully balanced the pole at the origin for 125 episodes, demonstrating stable control using counterfactual planning.
- After training, the agent could follow a time-varying target position $x(t) = A\sin(2\pi t/T)$ for multiple oscillation periods, even though this task was not part of the training objective.
- The agent maintained balance for up to 5000 steps in trials with a moving target, with average hold times exceeding one full period of oscillation for a range of amplitudes and periods.
- Performance degraded for very short periods (T < 250), where the agent could not track the target, indicating a limit on temporal tracking speed.
- The latent space dimension $N_L$ was the most critical hyperparameter; too large a space led to failure in balancing due to exploitation of unrealistic future predictions.
- The method achieved task success in fewer than 100 episodes on average, with the fastest solution found in around 35 episodes, comparable to state-of-the-art DQN and DQN-based agents.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.