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.