Neural Network Architectures for RL

Selecting the right neural network architecture is fundamental when implementing deep reinforcement learning agents. Just as the foundation of a building dictates its stability and height, the network structure profoundly influences the agent's learning capacity, stability during training, and overall performance. While the core RL algorithms define how an agent learns, the network architecture defines what it can represent and how efficiently it processes information from the environment.

This section focuses on practical considerations for designing these networks, tailored to the specific demands of RL tasks, building upon your existing knowledge of deep learning components.

Processing Environmental Observations

The first step in designing your network is handling the input: the observations ($o_t$) provided by the environment. The nature of the observation space largely dictates the initial layers of your network.

• Vector Observations: For environments where the state is represented as a fixed-size vector of numerical features (e.g., cart-pole positions and velocities, robot joint angles), standard Multi-Layer Perceptrons (MLPs) are usually effective. These typically consist of several fully connected (dense) layers with non-linear activation functions like ReLU. Input feature scaling (normalization or standardization) is often very important for stable training with MLP inputs.

• Image Observations: When the agent perceives the environment through images (e.g., pixels from game screens like Atari, or camera feeds in robotics), Convolutional Neural Networks (CNNs) are the standard choice. A common pattern, inspired by computer vision successes, involves stacking several convolutional layers (often with ReLU activations and possibly pooling layers) to extract spatial hierarchies of features, followed by one or more fully connected layers to produce the final output (Q-values or policy parameters).

  • Preprocessing: Raw pixels are rarely used directly. Common preprocessing steps include converting images to grayscale, downsampling, and stacking consecutive frames (e.g., 4 frames) to provide the network with information about motion and dynamics.
  • Example Structure: Input (stacked frames) -> Conv2D -> ReLU -> Conv2D -> ReLU -> Conv2D -> ReLU -> Flatten -> Dense -> ReLU -> Output Layer.

• Sequential Observations: If observations have inherent temporal dependencies that are not fully captured by simple frame stacking, or if the input is naturally sequential (e.g., text in text-based games, time-series data), Recurrent Neural Networks (RNNs) like LSTMs (Long Short-Term Memory) or GRUs (Gated Recurrent Units) can be employed. These networks maintain an internal hidden state, allowing them to integrate information over time. Using RNNs adds complexity to training, requiring techniques like Backpropagation Through Time (BPTT) and careful handling of hidden states, especially when using experience replay.

• Transformer Architectures: For tasks involving very long-range dependencies or structured data (like natural language instructions or complex relational environments), Transformer networks, utilizing self-attention mechanisms, are becoming increasingly viable. They excel at modeling dependencies across distant points in a sequence but come with higher computational costs and data requirements compared to CNNs or RNNs.

Generating Agent Actions

The output layer(s) of the network must match the action space ($a_t$) of the environment.

• Discrete Action Spaces: For environments with a finite set of actions (e.g., 'left', 'right', 'fire'), the network typically outputs a vector where each element corresponds to an action.

  • Value-Based Methods (e.g., DQN): The output layer provides the estimated Q-value $Q(s, a)$ for each possible action $a$ in state $s$. A linear output layer is common.
  • Policy-Based Methods (e.g., A2C, PPO): The output layer provides logits (unnormalized log-probabilities) for each action. These logits are then typically passed through a Softmax function to form a categorical probability distribution $\pi(a|s)$ from which actions are sampled.

• Continuous Action Spaces: For environments where actions are continuous vectors (e.g., motor torques, steering angles), the network outputs parameters defining a probability distribution over actions.

  • Deterministic Policies (e.g., DDPG): The network directly outputs the action vector $a = \mu(s)$. An activation function like tanh is often applied to the output layer to bound the actions within a specific range (e.g., [-1, 1]), which can then be scaled to the environment's requirements.
  • Stochastic Policies (e.g., PPO, SAC): The network typically outputs the parameters of a distribution, most commonly the mean $\mu(s)$ and standard deviation $\sigma(s)$ (or log standard deviation for stability) of a Gaussian distribution. Actions are then sampled from $\mathcal{N}(\mu(s), \sigma(s)^2)$. Care must be taken to ensure the standard deviation remains positive (e.g., by using a softplus activation or exponentiating the network output). The critic network in these Actor-Critic setups still outputs a scalar value estimate ($V(s)$ or $Q(s, a)$).

Common Network Layouts for Actor-Critic Methods

In Actor-Critic algorithms, you need networks for both the policy (actor) and the value function (critic). A common design choice is whether to share parameters between them.

• Separate Networks: Use entirely independent networks for the actor and the critic. This is simpler to implement and debug, and avoids potential negative interference between the policy and value function learning objectives. However, it might be less sample efficient as features relevant to both tasks are learned separately.

• Shared Backbone with Separate Heads: Use common initial layers (e.g., CNN layers for image processing) to extract a shared feature representation from the observation. Then, attach separate output layers ("heads") for the actor (policy output) and the critic (value output). This encourages learning shared representations, potentially improving sample efficiency and generalization. However, balancing the gradients and learning rates for the shared and separate parts can be challenging.

Comparison between separate networks and a shared backbone architecture for Actor-Critic methods.

The Dueling Architecture (Value-Based)

As introduced in Chapter 2, the Dueling Network Architecture is a specific modification primarily for value-based methods like DQN. It splits the network into two streams after the feature extraction layers:

1. Value Stream: Estimates the state value function $V(s)$. Outputs a single scalar.
2. Advantage Stream: Estimates the advantage function $A(s, a)$ for each action. Outputs a vector with dimensionality equal to the number of actions.

These streams are then combined to produce the final Q-values. A common aggregation method is:

$$Q(s, a) = V(s) + \left( A(s, a) - \frac{1}{|\mathcal{A}|} \sum_{a'} A(s, a') \right)$$

This centering ensures identifiability and improves stability. The intuition is that separating the state's intrinsic value from the action-specific advantages can lead to more effective learning, especially when the value of being in a state doesn't strongly depend on the specific action taken.

Structure of a Dueling Network Architecture, separating value and advantage estimation before combining them for Q-values.

Architectural Design Principles and Tips

While there's no single "best" architecture for all RL problems, some guiding principles help:

• Start Simple: Begin with known, relatively shallow architectures that have worked for similar problems (e.g., the standard Atari CNN for DQN). Increase complexity (depth, width, recurrence) only if performance plateaus and you suspect representational capacity is the bottleneck.
• Input/Output Matching: Ensure the input layers match the observation space and output layers match the action space and algorithm requirements (Q-values, policy distribution parameters).
• Normalization: Input observation normalization is often important. Techniques like running mean/standard deviation filters or batch normalization can be applied, but use techniques like Layer Normalization or be cautious with Batch Normalization in RL due to potentially non-stationary data distributions from evolving policies and experience replay buffers. Reward scaling or clipping is also a common practice.
• Activation Functions: ReLU remains a common choice for hidden layers due to its simplicity and effectiveness in mitigating vanishing gradients. tanh is often used for bounding continuous actions. Choose activations appropriate for the specific output required (e.g., linear for Q-values, softplus for standard deviations).
• Initialization: Standard weight initialization techniques (e.g., orthogonal initialization, Xavier, He) are important starting points. Sometimes, initializing the final policy layer with very small weights is beneficial to encourage initial exploration.
• Consult Literature: Leverage existing research. Papers often specify the architectures used for benchmark tasks. Starting with a proven architecture significantly increases the chance of successful implementation.

Choosing and refining the network architecture is an iterative process, often requiring experimentation alongside hyperparameter tuning (covered next) to achieve optimal performance for a given RL task.