FNO-RL

Repository for RL algorithms with FNO (Fourier Neural Operator) Encoder

This repository contains implementations of various reinforcement learning algorithms enhanced with FNO encoders for different environments.

Prerequisites

Common Dependencies

pip install torch numpy wandb tqdm matplotlib

Environment-Specific Dependencies

For Carla Lane PPO:

pip install gym d4rl

For Rainbow (Atari):

pip install atari-py opencv-python

For State DMC SAC:

pip install hydra-core dmc2gym imageio termcolor tensorboard

For Offline RBVE:

pip install gym[atari]

Projects

1. Carla Lane PPO

Three variants of PPO implementation for Carla lane-keeping task:

PPO-pytorch (Baseline CNN)

cd Carla-Lane-PPO/PPO-pytorch
python main.py

PPO-pytorch-FNO (with FNO encoder)

cd Carla-Lane-PPO/PPO-pytorch-FNO
python main.py <seed>
# Example: python main.py 2024

Or use the provided script:

./run.sh

PPO-pytorch-FNO-proj (FNO with projection)

cd Carla-Lane-PPO/PPO-pytorch-FNO-proj
python main.py <seed>
# Example: python main.py 2024

Or use the provided script:

./run.sh

Note:

• Requires carla-lane-v0 environment from d4rl
• Logs to wandb project 'CARLA_PPO'
• Uses 4-frame stacking with grayscale conversion

2. Atari Efficient Rainbow

Rainbow DQN implementation with FNO encoder for Atari games.

cd Rainbow
python rainbow.py --game <game_name> --architecture fno-model [options]

Key arguments:

• --game: Atari game name (e.g., 'space_invaders', 'breakout')
• --architecture: Choose from 'canonical', 'data-efficient', or 'fno-model'
• --seed: Random seed (default: 2025)
• --T-max: Training steps (default: 50M)
• --gpu-id: GPU ID to use

Example:

python rainbow.py --game space_invaders --architecture fno-model --seed 2025

Note:

• Logs to wandb automatically
• Saves models and memory periodically

3. State DMC SAC

Soft Actor-Critic (SAC) with FNO encoder for DeepMind Control Suite tasks.

cd State-DMC-SAC
python train.py env=<env_name> [options]

Available environments:

• cartpole_swingup
• cheetah_run
• ball_in_cup_catch
• And other DMC tasks

Example:

python train.py env=cartpole_swingup seed=1

Or use the provided script:

./run.sh

Configuration:

• Edit config/train.yaml for training parameters
• Edit config/agent/sac.yaml for agent-specific parameters
• Logs to wandb project 'state_based_fno'

4. Offline RBVE (Regularized Bellman Value Estimator)

Offline RL with FNO encoder for Atari games using pre-collected datasets.

cd Offline-RBVE
python main_lr.py

Configuration: Edit the ArgumentStorage dictionary in main_lr.py:

• env_name: Atari environment (e.g., "Boxing-v0")
• data_path: Path to pre-collected dataset
• learning_rate: Learning rate (default: 1e-4)
• num_iterations: Number of training iterations

Note:

• Requires pre-processed offline Atari dataset
• Runs multiple learning rate experiments in parallel using threading
• Saves models and metrics automatically

5. ViZDoom A2C

A2C agent implementation for ViZDoom environments.

Available agents:

• a2c_agent.py - Advantage Actor-Critic
• dqn_agent.py - Deep Q-Network
• ddqn_agent.py - Double DQN
• ppo_agent.py - Proximal Policy Optimization
• reinforce_agent.py - REINFORCE algorithm

Usage: Import the desired agent from algos.agents and integrate with your training loop.

from algos.agents.a2c_agent import A2CAgent
from algos.preprocessing.stack_frame import StackFrame

# Initialize agent
agent = A2CAgent(input_shape, action_size, seed, device, 
                gamma, alpha, beta, update_every, actor_m, critic_m)

# Training loop
for episode in episodes:
    state = env.reset()
    # ... training logic

FNO Encoder

Each project includes an fno.py file with Fourier Neural Operator implementations:

• FNO1d - For 1D state spaces (State-DMC-SAC)
• FNO2d - For 2D image observations (Carla, Rainbow, Offline-RBVE)

The FNO encoder provides efficient spectral feature extraction for RL tasks.

Logging

All projects support Weights & Biases (wandb) logging. Configure your wandb entity in the respective training scripts.

Citation

If you use this code, please cite the relevant papers for the algorithms and FNO encoder.