Source-linked AI summary
Schema Networks: Zero-shot Transfer with a Generative Causal Model of Intuitive Physics
Ken Kansky, Tom Silver, David A. Mély, Mohamed Eldawy, Miguel Lázaro-Gredilla, Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, Dileep George
TL;DR
Task-to-task transfer remains limited despite strong progress on individual reinforcement-learning tasks. The paper introduces Schema Networks, an object-oriented generative model that learns causal dynamics and supports planning, then evaluates them on Breakout variations against A3C and Progressive Networks. Schema Networks achieve efficient, robust generalization and transfer, including zero-shot performance, while the experiments remain limited to simplified deterministic binary environments.
Problem
Task-to-task transfer remains limited despite substantial progress on individual reinforcement-learning tasks, motivating models that generalize across related environments.
Method
Schema Networks are object-oriented generative models that learn entity-based causal dynamics and use probabilistic inference for action selection and planning.
Results
Schema Networks show more efficient and robust generalization than A3C and Progressive Networks across a suite of Breakout variations, including zero-shot transfer.
Takeaways & Limitations
Explicit causal structures support transfer to novel tasks and enable planning from goals through causal relationships.
Takeaways & Limitations
The experiments use binary states, actions, and rewards in deterministic environments with no uncertainty in observed entity states.
Abstract
from arXiv · showhide
The recent adaptation of deep neural network-based methods to reinforcement learning and planning domains has yielded remarkable progress on individual tasks. Nonetheless, progress on task-to-task transfer remains limited. In pursuit of efficient and robust generalization, we introduce the Schema Network, an object-oriented generative physics simulator capable of disentangling multiple causes of events and reasoning backward through causes to achieve goals. The richly structured architecture of the Schema Network can learn the dynamics of an environment directly from data. We compare Schema Networks with Asynchronous Advantage Actor-Critic and Progressive Networks on a suite of Breakout variations, reporting results on training efficiency and zero-shot generalization, consistently demonstrating faster, more robust learning and better transfer. We argue that generalizing from limited data and learning causal relationships are essential abilities on the path toward generally intelligent systems.
1. Introduction
Task-to-task transfer remains limited despite strong single-task progress, especially when Breakout object positions change while dynamics and rewards remain constant. Schema Networks are introduced to model causal structure and support robust generalization across such variations.
- Motivation: Deep reinforcement learning models often struggle to transfer experience from one environment to another.The limitation is illustrated by Breakout variations that perturb object positions without changing movements or reward sources.
- Evaluation: After training on standard Breakout, Schema Networks generalize to several positional and rule variations without additional training.The variations include a middle wall, altered brick rewards, an offset paddle, random targets, and juggling.
- Motivation: Object-based representations capture entities and interactions but do not by themselves model causality or explain alternative causes.The paper links causal modeling to regression planning, which reasons backward from desired future states.
- Contribution: Schema Networks represent local cause-effect relationships among entities and use them to guide action selection in new settings.Their representation is designed to handle uncertainty, multiple causation, and explaining away.
- Evaluation: Schema Networks are evaluated against A3C and Progressive Networks on Breakout variations for training efficiency and transfer.The evaluation targets generalization beyond the training environment rather than performance on a single task.
2. Related Work
Related work improves reinforcement learning through asynchronous experience, object-centered structure, relational modeling, and transfer architectures. Schema Networks differ by combining structured causal representations with generative inference and zero-shot transfer.
- Deep reinforcement learning: A3C accelerates reinforcement learning by combining experience from multiple agents acting asynchronously in copies of the same environment.The paper uses A3C as its primary comparison method.
- Object and relational models: Interaction Networks and Neural Physics Engines learn object-level and pairwise relational models of intuitive physics.Their primary advantage is compatibility with gradient-based learning methods.
- Transfer methods: Progressive Networks transfer features by freezing previously trained A3C copies and adding lateral connections to a new network for each task.This architecture still requires successive task-specific training.
- Schema Networks: Schema Networks transfer across scenarios with repeatable structure and sub-structure, while related model-free methods may require continued learning on the test environment.The paper identifies zero-shot transfer as a distinguishing capability.
- Probabilistic models: Probabilistic graphical models support inference with uncertainty and explaining away, making them useful for reinforcement learning and planning.The paper implements Schema Networks as probabilistic graphical models.
3. Schema Networks
Schema Networks are structured generative models that represent environments as entities with binary attributes and learn reusable schemas for transitions and rewards. Grounded schema factors, self-transitions, and probabilistic inference combine to model multiple causes and support planning.
- Model representation: A Schema Network parses an image into entities whose attributes form binary variables describing entity and complete environment states.The model assumes entities share a collection of attributes, with states represented as assignments to those attributes.
- Schemas: Grounded schemas predict future entity-attributes or rewards from present entity-attributes and optional actions when their preconditions are satisfied.Multiple grounded schemas can predict the same attribute, and their outputs are combined with an OR operation.
- Schemas: Ungrounded schemas act as templates that are instantiated across times, positions, and combinations of entities to create grounded schemas.This allows one learned cause-effect pattern to apply to multiple specific entities and locations.
- Transitions: Self-transition variables preserve entity-attributes when no active schema predicts a change.They also allow continuous or categorical attributes to be approximated with sets of binary variables.
- Planning: The factor graph represents transitions and rewards probabilistically, enabling planning as inference over action sequences.The explicit OR of possible causes supports reasoning about multiple causation and regression planning.
- Assumptions: The model relies on a vision system to detect and track entities and can use observable or latent entity attributes.The paper restricts its experiments to fully observable attributes, while noting possible latent attributes such as bounciness or magnetism.
4. Learning and Planning in Schema Networks
Schema Networks learn dynamics from entity-state sequences by inducing Boolean schemas, then use probabilistic inference and backtracking to plan actions toward reward goals.
- 4.1. Training Procedure: During data collection, planning selects actions with ε-greedy exploration, and accurate dynamics predictions eliminate the need for additional policy learning or policy adaptation.Transfer is therefore based on learned game dynamics rather than retraining a task-specific policy.
- 4.1. Training Procedure: Training converts temporal entity states into binary input-output examples, then learns schemas that predict entity attributes from neighboring attributes.The representation uses a binary matrix X of local entity and neighbor attributes and a binary target vector y.
- 4.2. Schema Learning: Schema structure is learned greedily by adding high-precision schemas through LP relaxations, while minimizing prediction error subject to a complexity trade-off.The exact objective is NP-hard, so the method iteratively adds schema columns that increase recall without sacrificing precision.
- 4.2. Schema Learning: The learned model represents local cause-effect relationships with binary variables, Boolean compatibility factors, and reusable ungrounded schema templates.Ungrounded schemas instantiate across times and entity combinations, while grounded schemas predict attributes or rewards from satisfied preconditions.
- 4.3. Planning as Probabilistic Inference: Planning first estimates potentially feasible desirable states, clamps selected reward variables, and solves for action assignments using MAP inference.The procedure combines forward max-product belief propagation with backward search because the factor graph is loopy and the forward pass can be overoptimistic.
- 4.3. Planning as Probabilistic Inference: Backtracking verifies feasibility and finds satisfying configurations, handling positive rewards, avoidable negative rewards, and mutually exclusive causal alternatives.The planner combines a Viterbi-like backward pass with depth-first search; a single forward pass can incorrectly treat incompatible causes as jointly feasible.
5. Experiments
The experiments compare Schema Networks with A3C and Progressive Networks across Breakout variations that preserve dynamics while changing layouts or reward arrangements. Schema Networks transfer learned causal dynamics more effectively, including zero-shot performance and identifying reward causes.
- 5. Experiments: The Breakout variations share dynamics but alter layouts, requiring different policies and diverse concepts for predicting object movements and rewards.The experiments test whether models can separate relevant causes, such as brick color, from irrelevant factors such as ball velocity and hit position.
- 5. Experiments: Schema Networks receive entity states, whereas A3C and Progressive Networks receive augmented images containing object-shape and paddle-position channels.Schema Networks must also learn which of their available object-part attributes are irrelevant for prediction.
- 5.2. Zero-Shot Generalization: Schema Networks perform zero-shot transfer to the Middle Wall variation after pretraining because its required dynamics schemas were already learned.Progressive Networks require substantial additional training in the new environment, while Schema Networks ignore the extra data.
- 5.2. Zero-Shot Generalization: In the zero-shot score comparison, Schema Network averages exceed the best A3C scores by more than one standard deviation in every variation.A3C used the two best of five attempts after 200k Standard Breakout frames, while Schema Networks averaged five attempts after 100k Mini Breakout frames.
- 5.2. Zero-Shot Generalization: Schema Networks outperform A3C in every tested Breakout variation, including Offset Paddle, Middle Wall, Random Target, and Juggling.The authors attribute these differences to Schema Networks transferring causal dynamics for aiming and responding to imminent negative rewards, whereas A3C struggles to adapt its policy.
- 5.3. Testing for Learned Causes: Schema Networks learn which brick colors cause positive or negative rewards and prefer the positive half at test time, while A3C achieves an average score near zero.Training used random negative-brick positions; testing grouped positive and negative colors into separate halves.
6. Discussion and Conclusion
Schema Networks perform strongly across Breakout variations, learning causality rather than directly maximizing rewards; their simplified experimental environments leave continuous, stochastic settings for future work.
- Schema Networks achieve rewards more efficiently than model-free methods such as A3C while learning an accurate game model.High scores are described as a byproduct of modeling game causality.
- Their success is partly attributed to representing state through entities, which may improve Deep RL training efficiency and generalization.
- The experiments discretize states, actions, and rewards as binary variables, use deterministic dynamics, and assume no uncertainty in observed entity states.
- Schema Networks show promise for multi-task transfer where Deep RL struggles, using explicit causal structures to support planning in novel tasks.
A. Breakout playing visualizations
Visualizations show typical gameplay in the Midwall Breakout variation, where Schema Networks usually avoid the middle wall while A3C struggles to score.
- Figure 4 presents typical gameplay for one Breakout variation.
- In the Midwall variation, Schema Networks purposefully avoid the middle wall most of the time, whereas A3C struggles to score points.
B. LP-based Greedy Schema Learning
LP-based greedy schema learning incrementally finds schemas for solvable input clusters, sparsifies them while preserving precision and recall, and adds the resulting schema to the network.
- Algorithm 1 specifies the LP-based greedy schema-learning procedure.
- LP-based greedy schema learning finds input-sample clusters solvable by one relaxed schema while maintaining perfect precision.The procedure selects an input sample, marks it solved, and solves a linear program.
- The resulting schema is simplified by marking zero-output samples solved and making its weights as sparse as possible without changing precision or recall.
- The learned schema is binarized when necessary and then added to the network as a new schema.