We introduce JAXbounce, a continuous control pong-like benchmark designed to evaluate active inference beyond discrete, low dimensional environments. JAXbounce combines continuous states and stochastic dynamics in discrete time with flexible and composable observation models that allow representing the environment as a state vector or a pixel grid, corruptible through the rich family of noise models we provide. It is fully differentiable and parallelisable in JAX. We run over 10,000 agent simulations on our environment comparing a baseline deterministic solution, to a state-based active inference agent with Unscented Kalman Filter perception and expected free energy planning, and a pixel based agent whose measurement model is a pretrained neural network encoder. All agents are evaluated across three scenarios of increasing difficulty with eight different noise conditions. Our experiments characterise how performance of continuous state agents scales with planning, planning horizons, observation modalities and knowledge of the noise process, exhibiting preference-oriented behaviour even in the presence of increasing uncertainty. We release JAXbounce as a reproducible benchmark for continuous control active inference research.
Figures & tables
Figure 1: Experiment 5. State AIF performance heatmaps across all sampled plan variations ( N={100,200,300,400} ) and planning horizons ( H={20,60,100} ). The agent here has an initial belief covariance of the noise similar to the actual noise conditions listed in table 1 .
Figure 2: Sensitivity analysis of the agent’s chosen plan for four example noise conditions by minimizing the NEFE with respect to all actions in its planning horizon. The plot shows the partial derivative of the actions in the plan with respect to the NEFE. An increase in the number of sampled plans decreases the gradient such that a small perturbation in the actions does not change the agent’s chosen plan’s NEFE.
Figure 3: Sensitivity analysis of the agents planning for four example noise conditions across different number of plans. The plot shows the state AIF agent’s sensitivity to the plan chosen by minimizing the NEFE to all the actions in its planning horizon for the case where the agent’s belief of the noise matches the actual noise. An increase in the number of sampled plans decreases the gradient.
Figure 4: (Left) Evolution of the state AIF agent’s belief about the friction field around the bat vs the true friction field. (Middle) The agent’s surprise over its observation vector (Right) Agent’s surprise over balls’ x and y position and the bat’s lateral position.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Visualisations of the JAXbounce grids used in pixel AIF simulations. The clean grid with increasing uncertainty on the agent’s initial belief (table 1 ) of the position of the ball and the bat and increasing uncertainty on its measurement model is used for the deterministic and state AIF simulations.
Label
σpos(0)
σvmult
σoinj
Noise 1
10−3
1.0
10−3
Noise 2
0.02
1.5
0.05
Noise 3
0.05
2.0
0.10
Noise 4
0.10
2.5
0.20
Noise 5
0.12
3.0
0.25
Noise 6
0.15
3.5
0.30
Appendix
Table 1: Noise conditions in experiments for the deterministic baseline and state-based AIF agent. The eight noise tiers are designed to make inference increasingly harder for the agents.
Figure 6: Experiment 1. Bat size comparison heatmaps between the deterministic and state AIF agents. All runs are averaged over 30 trials for Scenario 1 (constant velocity ball falling), Scenario 2 (variable ball velocity in each run) and Scenario 3 (variable ball velocity and lateral ball falling). State AIF agent runs randomly sample 300 plans over a finite horizon of 100 timesteps.
Figure 7: Experiment 2. Analysis of agent performance for varying number of plans across the three scenarios. All plans are run for a horizon of 100.
Figure 8: Experiment 3. Comparison of performance for varying planning horizons across the three scenarios. The agents randomly sample 300 plans.
Figure 9: Experiment 4. Varying grid sizes for the pixel AIF agent. Each simulation is run for 300 randomly sampled plans over a planning horizon of 100 time steps.
Sophisticated Inference is a variant of active inference often associated with recursive belief modeling and tree search. We argue that its central computational role is simpler: within a planning horizon, it makes active inference closed-loop by allowing future actions to depend on future states and observations. This closed-loop structure can be represented in the epistemic-prior variational free energy framework. Epistemic priors supply the active-inference objective, while a joint posterior over future states and actions supplies the state-contingent control structure. We evaluate this decomposition in the Reactivity Maze, a stochastic benchmark designed to separate epistemic incentive from inner-horizon closed-loop control. The comparison includes three variational objectives with the same state-action posterior family, an action-state factorized active inference objective, Sophisticated Inference, and standard Expected Free Energy planning. The results show that neither ingredient is sufficient on its own. Methods without an epistemic component do not seek information, while methods that prevent future actions from depending on future states cannot turn information into reliable goal-reaching. By contrast, both Sophisticated Inference and full-joint epistemic-prior active inference solve the environment by combining epistemic drive with closed-loop inference. These results show that the advantage associated with Sophisticated Inference need not be specific to tree search itself. It arises from the closed-loop form of active inference, and this form can be represented in epistemic-prior variational inference when the posterior keeps future actions dependent on future states.
Wouter W. L. Nuijten, Bert de Vries
Eindhoven University of Technology, 5612 AP Eindhoven, the Netherlands · Lazy Dynamics, Utrecht, the Netherlands
Learning from trajectory demonstrations offers a route to active-inference control of complex systems whose dynamics are difficult to model explicitly. We introduce generative active-inference control (GenAIF), in which one generative trajectory model learns from demonstrations and measured action interventions to supply a goal-conditioned policy distribution and a state-to-observation likelihood mapping. From this control design, we derive three model requirements: (i) useful action proposals, (ii) accurate prediction under imposed actions, and (iii) probabilistic observation evidence for belief updating and expected information gain. We benchmark diffusion, autoregressive Transformers, conditional variational autoencoders (CVAEs), and flow matching in a MuJoCo manipulation task with multiple physical conditions. Diffusion delivers the strongest control across the tested dynamics, while CVAE combines comparable short-horizon prediction with much faster inference. Correct conditioning is decisive, and trajectory reuse offers further computational savings. In replay after an unannounced tilt change, pretrained diffusion updates belief fastest among the original models; fine-tuning on recovery demonstrations further accelerates identification and sustains accurate tracking. These findings support the use of shared generative trajectory models to connect action proposal, controlled prediction, and observation evidence within GenAIF.
Yulin Li, Mohsen A. Jafari, Andrea Matta
Rutgers University–New Brunswick, Piscataway NJ 08854, USA · Politecnico di Milano, Via La Masa 1, 20156 Milano, Italy
Active inference, a neurally-inspired model for inferring actions based on the free energy principle (FEP), has been proposed as a unifying framework for understanding perception, action, and learning in the brain. Active inference has previously been used to model ecologically important tasks such as navigation and planning, but scaling it to solve complex large-scale problems in real-world environments has remained a challenge. Inspired by the existence of multi-scale hierarchical representations in the brain, we propose a model for planning of actions based on hierarchical active inference. Our approach combines a hierarchical model of the environment with successor representations for efficient planning. We present results demonstrating (1) how lower-level successor representations can be used to learn higher-level abstract states, (2) how planning based on active inference at the lower-level can be used to bootstrap and learn higher-level abstract actions, and (3) how these learned higher-level abstract states and actions can facilitate efficient planning. We illustrate the performance of the approach on several planning and reinforcement learning (RL) problems including a variant of the well-known four rooms task, a key-based navigation task, a partially observable planning problem, the Mountain Car problem, and PointMaze, a family of navigation tasks with continuous state and action spaces. Our results represent, to our knowledge, the first application of learned hierarchical state and action abstractions to active inference in FEP-based theories of brain function.
Prashant Rangarajan, Rajesh P. N. Rao
Paul G. Allen Center for Computer Science and Engineering, University of Washington, Seattle, USA