Adapting Critic Match Loss Landscape Visualization to Off-policy Reinforcement Learning

1. Problem Statement

Reinforcement Learning (RL) algorithms, particularly Actor-Critic methods, are often treated as "black boxes," making it difficult to interpret their internal optimization dynamics and diagnose training instability. While loss landscape visualization has been successfully applied to online RL (where data is collected step-by-step and targets are updated immediately), it cannot be directly applied to off-policy RL (e.g., Soft Actor-Critic, SAC).

Off-policy RL introduces two structural complexities that break standard visualization methods:

1. Replay-based Data Flow: Updates occur on mini-batches sampled from a replay buffer rather than immediate step-wise interactions.
2. Target Computation: Target values (Bellman targets) rely on delayed target networks, twin-critic structures, and entropy regularization, making the loss function non-stationary and dependent on future policy states.

The core problem is the lack of a geometric diagnostic tool to visualize and analyze the critic optimization landscape in off-policy settings, which is crucial for understanding why certain algorithms converge while others diverge.

2. Methodology

The authors propose an adaptation of the Critic Match Loss Landscape visualization method to the Soft Actor-Critic (SAC) algorithm. The methodology involves three main components:

A. Adaptation of the Visualization Framework

To bridge the gap between online and off-policy RL, the authors modified the visualization procedure in two key areas:

1. Data Flow Adaptation:
   • Instead of using continuously evolving online samples, a fixed batch of replay data is sampled from the buffer.
   • Critic weights are recorded at the frequency of neural network update steps (rather than episode boundaries) to capture the rapid parameter changes inherent in off-policy training.
2. Loss Computation Adaptation:
   • Frozen Targets: To create a stationary loss surface, the Bellman target values ($y_t$) are precomputed and fixed based on the final policy. This eliminates the ambiguity caused by recomputing targets for every perturbed parameter point.
   • Twin-Critic Handling: The visualization focuses on a single primary critic network, as both critics share the same loss formulation and target definition.
   • Freezing Auxiliary Components: The actor network and the temperature coefficient ($\alpha$) are frozen to isolate the critic's local loss geometry.

B. Construction of the Landscape

• Dimensionality Reduction: High-dimensional critic weights recorded during training are projected onto a 2D plane using Principal Component Analysis (PCA). The two principal components ($\delta, \eta$) serve as the axes for the landscape.
• Loss Evaluation: For a grid of points $(\alpha, \beta)$ on this plane, the critic parameters are perturbed ($w' = w + \alpha\delta + \beta\eta$), and the Critic Match Loss is evaluated using the fixed replay batch and frozen targets.
• Visualization: This generates a 3D loss surface with a 2D optimization path overlaid, showing the trajectory of the critic parameters relative to the loss geometry.

C. Quantitative Metrics

To move beyond qualitative observation, three metrics are defined:

1. Sharpness ($Sharp_\epsilon$): Measures the worst-case sensitivity of the loss to perturbations within a radius $\epsilon$. High sharpness implies a well-defined but potentially fragile solution; low sharpness suggests a flat, potentially unstable region.
2. Basin Area ($A_\rho$): The area of the region where the normalized loss remains below a threshold $\rho$. A large basin indicates robustness to parameter perturbations.
3. Local Anisotropy ($\log \kappa$): The logarithm of the condition number of the Hessian matrix. High anisotropy indicates an ill-conditioned valley (steep in one direction, flat in another), often leading to slow convergence or instability.

3. Key Contributions

1. Methodological Adaptation: Successfully extended the critic match loss landscape visualization from online to off-policy RL by addressing replay buffers and target network delays.
2. Geometric Diagnostic Tool: Established a framework to visualize the "frozen-target" surrogate objective, allowing for the analysis of critic optimization dynamics independent of the non-stationary training process.
3. Comparative Analysis: Provided a geometric explanation for the differences between SAC (off-policy) and ADHDP (online), and between convergent and divergent training outcomes.
4. Temporal Snapshots: Introduced a method to visualize the evolution of the loss landscape over time, revealing how optimization trajectories shift or fragment during training.

4. Results

The method was validated on two control tasks: Spacecraft Attitude Control and the Cart-Pole benchmark.

A. Convergent SAC vs. Divergent ADHDP (Spacecraft)

• SAC (Convergent): Exhibited a broad, smooth, concave basin with a single valley. The optimization path was continuous and converged into a stable low-loss region.
  • Metrics: Large basin area ($A_\rho \approx 1252$), moderate anisotropy ($\log \kappa \approx 3.7$).
• ADHDP (Divergent): Showed sharp curvature, irregular contours, and no closed basin. The optimization path oscillated in a narrow, restricted parameter space.
  • Metrics: Non-existent basin, high anisotropy ($\log \kappa \approx 10.4$).
• Insight: The smooth geometry of SAC correlates with its stability, while the rugged, anisotropic geometry of ADHDP reflects sensitivity to noise and unstable target updates.

B. Divergent SAC Cases (Spacecraft)

The authors analyzed two types of SAC divergence to show that geometric stability does not always guarantee control success:

1. Divergent Parameters (Runaway): The critic weights grew unbounded. The landscape showed strong anisotropy (effectively 1D), where the trajectory was driven along a single dominant direction (ridge) rather than into a basin.
2. Convergent Parameters (Discontinuous): The critic weights appeared to converge, but the control failed. The landscape showed a smooth surface, but the optimization trajectory was fragmented, jumping between disconnected clusters. This revealed that the failure was due to discontinuous regime shifts in the optimization path, not the shape of the loss surface itself.

C. Cart-Pole Validation

The method was successfully applied to the Cart-Pole system, yielding results consistent with the spacecraft case: convergent SAC produced a smooth, well-defined basin, confirming the method's generalizability across different dynamical systems.

5. Significance

• Diagnostic Power: The framework serves as a powerful geometric diagnostic tool. It can detect instability mechanisms (e.g., directional runaway, trajectory fragmentation) that are invisible in standard scalar training curves (loss/reward plots).
• Algorithmic Insight: It clarifies why off-policy algorithms like SAC are generally more robust than online methods like ADHDP, attributing this to the formation of smoother, wider loss basins facilitated by replay buffers and target smoothing.
• Limitations & Scope: The authors clarify that the visualization represents a stationary surrogate objective (frozen targets). While it reveals structural properties of the critic, it does not perfectly reconstruct the full non-stationary RL objective. Therefore, a stable landscape does not guarantee closed-loop control success if the actor-critic coupling introduces other instabilities.
• Future Impact: This work provides a standardized approach for researchers to visualize and interpret the internal dynamics of modern off-policy RL algorithms, aiding in hyperparameter tuning and the development of more stable control policies.