Adaptive Correlation-Weighted Intrinsic Rewards for Reinforcement Learning

This paper proposes ACWI, an adaptive intrinsic reward framework that employs a lightweight Beta Network to dynamically learn state-dependent scaling coefficients via a correlation-based objective, thereby enhancing sample efficiency and training stability in sparse reward reinforcement learning environments without the need for manual tuning.

The Gist

Imagine you are teaching a robot to navigate a giant, dark maze to find a hidden treasure. The problem is that the robot only gets a "ding!" of happiness when it actually finds the treasure. For the first 99% of the journey, it gets absolutely no feedback. It's like walking in the dark and hoping you don't bump into a wall, but you have no idea if you're getting closer to the exit or just walking in circles.

This is the classic problem of Reinforcement Learning in "sparse reward" environments. The robot needs to explore, but without a guide, it often gets lost or gives up.

The Old Way: The "One-Size-Fits-All" Compass

To help the robot, scientists previously gave it a "curiosity compass." This compass gave the robot a little bonus point every time it visited a new, unknown spot in the maze. This encouraged it to wander around and explore.

However, there was a catch. Scientists had to manually decide how strong this curiosity should be.

• If they set it too low, the robot stayed lazy and never explored enough.
• If they set it too high, the robot got too curious. It would run around frantically, ignoring the actual path to the treasure just to see what was behind the next door.

It was like trying to tune a radio with a dial that you could only set once at the beginning of the day. If the signal was clear in the morning but noisy in the afternoon, the radio would either be too quiet or full of static the whole time. You needed a way to adjust the volume while you were listening.

The New Solution: ACWI (The Smart Volume Knob)

The paper introduces a new method called ACWI (Adaptive Correlation-Weighted Intrinsic). Think of ACWI as a smart, self-adjusting volume knob for the robot's curiosity.

Instead of a fixed setting, ACWI uses a tiny, lightweight "brain" (called a Beta Network) that looks at the robot's current situation and asks: "Is being curious right now going to help me find the treasure?"

Here is how it works using a simple analogy:

1. The Two Types of Rewards

• Extrinsic Reward (The Treasure): The actual goal (finding the key, opening the door, reaching the goal). This is the "real" money.
• Intrinsic Reward (The Curiosity): The "fun" of discovering something new. This is the "allowance" the robot gets for exploring.

2. The Beta Network (The Smart Manager)

Imagine the robot is a student taking a test.

• The Old Way: The teacher says, "You get 10 bonus points for every new page you read," regardless of whether the page is useful or just random scribbles.
• The ACWI Way: The teacher has a Smart Manager watching the student.
  • If the student is in a section of the book where reading leads to the answer key, the Manager says, "Great! Read more! Turn up the curiosity volume!" (High bonus points).
  • If the student is wandering into a section that is just a dead-end or a wall, the Manager says, "Stop wasting time. Turn down the curiosity volume." (Low bonus points).

3. How Does the Manager Know? (The Correlation Trick)

You might ask, "How does the manager know which path leads to the treasure if the robot hasn't found it yet?"

The manager looks at the future. It uses a clever trick called Correlation.

• It watches the robot's actions.
• It asks: "When the robot was curious in this specific spot, did it eventually lead to a big reward later on?"
• If Yes: The manager learns to give a high "curiosity bonus" for that type of spot in the future.
• If No: The manager learns to ignore curiosity in that spot.

It's like a detective looking at a map. If every time the detective took a left turn at the bakery, they eventually found the suspect, the detective learns: "Okay, curiosity at the bakery is valuable. Let's reward that." If taking a left turn at the park led nowhere, the detective stops rewarding that path.

Why This is a Big Deal

1. No More Guessing: Scientists don't have to spend weeks manually tuning the "curiosity dial" for every new game or maze. The system figures it out on its own.
2. Stability: It prevents the robot from going crazy with curiosity when it doesn't need to, and keeps it curious when it does.
3. Graceful Failure: If the maze is so empty that there are no clues at all (like a completely blank room), the system realizes, "Hey, curiosity isn't helping right now," and it just acts like a normal, fixed system. It doesn't crash; it just adapts to the lack of information.

The Result

In their experiments, this "Smart Volume Knob" (ACWI) helped robots learn faster and more reliably than the old "One-Size-Fits-All" methods. The robots found the treasure more efficiently because they knew exactly when to be curious and when to focus on the goal.

In short: ACWI teaches the robot to be smart about its own curiosity, turning it up when it helps and turning it down when it doesn't, all without needing a human to constantly adjust the settings.

Technical Summary

1. Problem Statement

Reinforcement Learning (RL) agents struggle in sparse reward environments where extrinsic feedback is rare, making it difficult to distinguish productive behaviors from random exploration. While Intrinsic Motivation (IM) methods (e.g., ICM, RND) address this by adding curiosity-driven rewards based on prediction errors or novelty, they suffer from a critical limitation: static reward scaling.

Current approaches combine extrinsic rewards ($R^E_t$) and intrinsic rewards ($R^I_t$) using a fixed scalar coefficient ($\beta$):
$$\bar{r}_t = R^E_t + \beta R^I_t$$
This fixed $\beta$ is typically tuned manually via hyperparameter search. The paper argues that this is suboptimal because:

• State Dependence: The utility of exploration varies by state. Some states require high exploration to reach high extrinsic returns, while others (where the path is clear) require exploitation. A uniform $\beta$ cannot distinguish between these scenarios.
• Instability: A fixed $\beta$ that works for one task or training phase may cause over-exploration (destabilizing training) or under-exploration (failing to find the goal) in another.
• Lack of Alignment: Fixed coefficients do not guarantee that intrinsic rewards align with future task success.

2. Methodology: ACWI

The authors propose ACWI (Adaptive Correlation-Weighted Intrinsic), a framework that learns a state-dependent scaling factor $\beta(s_t)$ online to dynamically balance intrinsic and extrinsic rewards.

Core Components

1. Beta Network ($\beta_\psi$):

   • A lightweight neural network (encoder + MLP head) that takes the agent's state $s_t$ as input and outputs a scalar scaling factor $\beta_\psi(s_t)$.
   • The output is constrained to a positive range $[\beta_{min}, \beta_{max}]$ (e.g., $[0.1, 2.0]$) via log-space clamping.
   • The augmented reward becomes: $\bar{r}_t = R^E_t + \alpha \cdot \beta_\psi(s_t) \cdot R^I_t$, where $\alpha$ is a global magnitude coefficient.

2. Intrinsic Module (ICM):

   • The authors instantiate the framework using the Intrinsic Curiosity Module (ICM).
   • ICM generates raw intrinsic rewards ($I_t$) based on the forward prediction error of a learned feature representation.
   • These rewards are standardized and rectified ($I^+_t$) to ensure stability.

3. Correlation-Based Training Objective:

   • Instead of using complex meta-learning or second-order gradients, ACWI optimizes the Beta Network using a first-order correlation objective.
   • Goal: Maximize the correlation between the weighted intrinsic signal ($\beta_\psi(s_t) \cdot I^+_t$) and the discounted future extrinsic return ($G^E_t$).
   • Mechanism: The network is trained to up-weight intrinsic rewards in states that lead to high future extrinsic returns and down-weight them in states where exploration is unproductive.
   • Loss Function:
     $$L_\beta(\psi) = -\mathbb{E}_B[\hat{I} \cdot \hat{G}] + \lambda_{reg} \mathbb{E}_B[(\log \beta_\psi(s) - \log \beta_0)^2]$$
     Where $\hat{I}$ and $\hat{G}$ are standardized (z-score) intrinsic and extrinsic returns within a mini-batch. The regularization term prevents the network from collapsing to extreme values.

4. Optimization Loop:

   • The Beta Network is updated via gradient descent once per training iteration, immediately before the PPO policy updates.
   • Crucially, the policy parameters ($\theta$) are detached during the Beta Network update, preventing second-order dependencies and ensuring computational efficiency.

3. Key Contributions

1. State-Dependent Scaling: Formulated intrinsic reward modulation as learning a state-dependent multiplier $\beta(s_t)$, allowing the agent to adapt exploration incentives granularly based on the current situation rather than global training stages.
2. Correlation Objective: Introduced a lightweight, first-order training objective that aligns intrinsic bonuses with future extrinsic returns without requiring unrolled meta-gradients or additional policy optimization loops.
3. Empirical Validation: Demonstrated that ACWI improves sample efficiency and learning stability across diverse sparse-reward environments compared to fixed-coefficient baselines, with minimal computational overhead.

4. Experimental Results

The authors evaluated ACWI on five MiniGrid environments (DoorKey-8x8, Empty-16x16, RedBlueDoors-8x8, UnlockPickup, KeyCorridorS3R3) using PPO.

• Performance: ACWI consistently outperformed PPO with fixed $\beta$ coefficients (tested across $\{0.1, 0.2, 0.5, 1, 2\}$) and standard PPO. It achieved faster early learning and reduced variance across random seeds.
• Adaptivity:
  • In structured environments (e.g., DoorKey), the learned $\beta$ distributions evolved from unimodal to multimodal, partitioning the state space into regions requiring high exploration vs. exploitation.
  • As training progressed and policies improved, $\beta$ values generally decreased, reflecting a natural shift from exploration to exploitation.
• Graceful Degradation: In the Empty-16x16 environment (extreme sparsity with almost no intermediate signals), the correlation signal collapsed. ACWI did not diverge; instead, the regularization term kept $\beta$ near its prior, effectively behaving like a fixed coefficient. This proved the method's robustness when correlation signals are uninformative.
• Visualization: PCA analysis showed that in structured tasks, $\beta$ values aligned with task-relevant geometric clusters in the state space, whereas in uninformative tasks, $\beta$ remained uniform.

5. Significance and Conclusion

• Practicality: ACWI offers a practical solution to the "tuning nightmare" of intrinsic rewards. It removes the need for manual hyperparameter search for $\beta$ and adapts automatically to different tasks and training phases.
• Efficiency: By using a simple correlation objective rather than complex meta-learning, ACWI integrates seamlessly with standard RL algorithms (like PPO) with negligible computational cost.
• Theoretical Insight: The work highlights that effective intrinsic motivation requires alignment with task goals. Exploration should be amplified only when it causally contributes to future extrinsic success.
• Limitations: The method relies on the existence of some extrinsic signal to drive the correlation. In completely barren environments with zero feedback until the very end, the adaptive mechanism defaults to a static baseline, though it remains stable.

In summary, ACWI represents a significant step forward in making intrinsic motivation robust and adaptive, moving beyond static heuristics toward dynamic, state-aware exploration strategies.