Active Inference Gym Agent

An implementation of an Active Inference agent for OpenAI Gym environments, based on the Free Energy Principle. This project demonstrates how biological principles of perception and action can be applied to reinforcement learning tasks.

🚀 Quick Start - Google Colab

New! Want to run this in your browser with no setup? Check out the Google Colab version:

• colab_notebook.py - Complete notebook with copy-paste ready code blocks
• COLAB_README.md - Detailed instructions for using in Colab
• Includes PPO comparison - Compare Active Inference with standard RL (PPO)

Just copy the code blocks into Google Colab cells and run! See COLAB_README.md for details.

---

What is Active Inference?

Active Inference is a theoretical framework from computational neuroscience that describes how agents interact with their environment. It's based on the Free Energy Principle, which proposes that biological systems minimize "free energy" - essentially, they minimize surprise and prediction errors.

Key Concepts

1. Generative Model: The agent maintains an internal model of how the world works, predicting future observations based on current state and actions.

2. Expected Free Energy: Actions are selected to minimize expected free energy, which balances:

   • Pragmatic value (exploitation): Achieve preferred/goal states
   • Epistemic value (exploration): Reduce uncertainty about the world

3. Perception and Action:

   • Perception updates beliefs about the current state
   • Action changes the world to match predictions

This differs from traditional RL where agents maximize reward. In active inference, agents fulfill predictions and resolve uncertainty.

Environment: CartPole-v1

This implementation uses the classic CartPole-v1 environment from Gymnasium (formerly OpenAI Gym):

• Objective: Balance a pole on a moving cart
• State Space: 4D continuous (cart position, cart velocity, pole angle, pole angular velocity)
• Action Space: 2 discrete actions (push left or right)
• Success: Keep the pole upright (within ±12°) for as long as possible

Architecture

The agent consists of two main neural network models:

1. Generative Model (World Model)

• Learns the environment dynamics: P(s_t+1 | s_t, a_t)
• Predicts next states given current state and action
• Estimates uncertainty about predictions

2. Preference Model (Value Model)

• Encodes preferred/desired states
• Similar to a value function in RL
• Trained using temporal difference learning

Action Selection

Actions are selected by minimizing expected free energy across all possible actions:

EFE(a) = -E[Pragmatic Value] + Epistemic Value
       = -Preference(s_next) + Uncertainty(s_next)

The agent uses softmax action selection to balance exploration and exploitation.

Installation

Requirements

• Python 3.8+
• PyTorch 2.0+
• Gymnasium
• NumPy
• Matplotlib

Setup

# Clone the repository
git clone <repository-url>
cd active_inference

# Install dependencies
pip install -r requirements.txt

Usage

Training an Agent

Run the training script:

python train.py

This will:

1. Train an active inference agent on CartPole-v1 for 500 episodes
2. Save the trained model to ./models/CartPole-v1_agent.pt
3. Generate training plots in ./models/CartPole-v1_training_results.png
4. Evaluate the trained agent for 10 episodes

Expected Results

After training, you should see:

• Episode rewards steadily increasing
• Episode length approaching the maximum (500 steps)
• Moving average reward converging around 400-500
• Generative model loss decreasing as the world model improves
• Preference model loss stabilizing as value predictions improve

Using a Trained Agent

from active_inference_agent import ActiveInferenceAgent
import gymnasium as gym

# Create environment
env = gym.make('CartPole-v1')

# Create agent
agent = ActiveInferenceAgent(
    state_dim=4,
    action_dim=2
)

# Load trained model
agent.load('./models/CartPole-v1_agent.pt')

# Run episode
state, _ = env.reset()
done = False
total_reward = 0

while not done:
    action = agent.select_action(state, explore=False)
    next_state, reward, terminated, truncated, _ = env.step(action)
    done = terminated or truncated
    total_reward += reward
    state = next_state

print(f"Total reward: {total_reward}")

Customization

Training Hyperparameters

You can modify hyperparameters in train.py:

agent = ActiveInferenceAgent(
    state_dim=state_dim,
    action_dim=action_dim,
    hidden_dim=128,          # Size of hidden layers
    learning_rate=3e-4,      # Learning rate for both models
    gamma=0.99,              # Discount factor
    tau=0.5,                 # Temperature for action selection (lower = more greedy)
    planning_horizon=5,      # Planning steps (future feature)
    buffer_size=10000        # Replay buffer size
)

Using Different Environments

To use a different Gym environment, simply change the environment name:

# For MountainCar
train_agent(env_name='MountainCar-v0', num_episodes=1000)

# For Acrobot
train_agent(env_name='Acrobot-v1', num_episodes=800)

# For LunarLander
train_agent(env_name='LunarLander-v2', num_episodes=1500)

Note: Continuous action spaces require modifications to the agent architecture.

Project Structure

active_inference/
├── active_inference_agent.py  # Core agent implementation
├── train.py                   # Training and evaluation script
├── requirements.txt           # Python dependencies
├── README.md                  # This file
└── models/                    # Saved models and plots (created during training)
    ├── CartPole-v1_agent.pt
    └── CartPole-v1_training_results.png

Technical Details

Free Energy Minimization

The agent minimizes variational free energy, which can be decomposed as:

F = Complexity - Accuracy
  = KL[q(s|o) || p(s)] - E_q[log p(o|s)]

Where:

• q(s|o) is the approximate posterior (belief about state)
• p(s) is the prior over states
• p(o|s) is the likelihood (generative model)

Expected Free Energy

For action selection, the agent minimizes expected free energy:

G(π) = E_q[log q(s) - log p(o,s|π)]
     = Risk + Ambiguity - Epistemic Value

Where:

• Risk: Expected distance from preferred states
• Ambiguity: Expected uncertainty about outcomes
• Epistemic Value: Expected information gain

References

1. Friston, K. (2010). The free-energy principle: a unified brain theory? Nature Reviews Neuroscience, 11(2), 127-138.

2. Friston, K. J., et al. (2015). Active inference and epistemic value. Cognitive neuroscience, 6(4), 187-214.

3. Da Costa, L., et al. (2020). Active inference on discrete state-spaces: A synthesis. Journal of Mathematical Psychology, 99, 102447.

4. Millidge, B., et al. (2021). Deep active inference agents using Monte-Carlo methods. arXiv preprint arXiv:2106.11423.

License

MIT License - feel free to use this code for your own projects!

Contributing

Contributions are welcome! Potential improvements:

• Implement continuous action spaces (using Gaussian policies)
• Add more sophisticated planning (tree search over future trajectories)
• Implement hierarchical active inference
• Add more environments and benchmarks
• Optimize performance for larger state/action spaces

Troubleshooting

Issue: Agent not learning / rewards stay low

• Try increasing the number of episodes (500 may not be enough for harder environments)
• Adjust learning rate (try 1e-3 or 1e-4)
• Increase hidden layer size for complex environments
• Check that the preference model is learning appropriate values

Issue: Training is slow

• Reduce replay buffer size
• Reduce batch size
• Use a GPU (agent automatically detects CUDA)
• Reduce network size (hidden_dim)

Issue: Agent is too exploratory/exploitative

• Adjust temperature parameter tau (lower = more greedy, higher = more exploratory)
• Modify the epistemic value weighting in the EFE calculation

Contact

For questions or issues, please open a GitHub issue or contact the maintainers.

---

Happy Active Inferencing! 🧠🤖