Shaping Sparse Rewards in Reinforcement Learning: A Semi-supervised Approach

Technical Summary: Shaping Sparse Rewards in Reinforcement Learning: A Semi-supervised Approach

Problem Statement
Reinforcement Learning (RL) faces a core challenge in real-world scenarios where reward signals are exceedingly sparse. In such environments, agents receive minimal immediate feedback, leading to high variance in returns and difficulty in learning optimal policies. While reward shaping and densification are common solutions, they often suffer from "reward hacking," where agents learn behaviors that deviate from the intended design to maximize artificial rewards. Conversely, supervised-based reward shaping methods, which rely on non-expert demonstrations, suffer from poor sample efficiency when non-zero-reward transitions are extremely rare. Existing approaches struggle to effectively utilize the vast majority of transitions that yield zero rewards.

Methodology: Semi-Supervised Reward Shaping (SSRS)
The authors propose the Semi-Supervised Reward Shaping (SSRS) framework, which bridges this gap by leveraging Semi-Supervised Learning (SSL) techniques. The core idea is to perform reward shaping not only on the scarce non-zero-reward transitions but also on the abundant zero-reward transitions by learning trajectory space representations.

The framework operates through the following key components:

1. Reward Estimation as Classification: The reward shaping problem is framed as a classification task. The system learns to assign a reward value from a discrete set $Z$ (containing observed non-zero rewards and zero) to state-action pairs. The reward estimator consists of two neural networks:

   • $R_V(s)$: A state-value reward distribution function mapping a state to a probability distribution over $Z$.
   • $R_Q(s, a)$: A state-action value reward distribution function mapping a state-action pair to a probability distribution over $Z$.
     The final confidence score vector $q_t$ for a transition is a weighted combination of these two distributions.

2. Semi-Supervised Optimization: The framework alternates between policy updates and reward estimator optimization.

   • Supervised Loss ($L_r$): For transitions with known non-zero rewards, the estimator is trained to minimize the discrepancy between its predictions and ground truth using a thresholded cross-entropy loss.
   • Consistency Regularization ($L_s$): For zero-reward transitions, the model employs consistency regularization. It applies weak and strong data augmentations to the input transitions. The model is optimized to produce consistent predictions between the weakly augmented and strongly augmented versions, provided the confidence exceeds a threshold $\lambda$. This utilizes the unlabeled zero-reward data to refine the reward function.

3. Entropy Augmentation: A novel data augmentation technique is introduced specifically for non-image tasks. Unlike traditional image augmentations (e.g., cutout, rotation) which may violate smoothness assumptions in trajectory space, this method utilizes Shannon Entropy. It calculates the entropy of the state matrix and uses it to perturb the state representation. This preserves the smoothness and clustering assumptions required for SSL to be effective in the trajectory space.

4. Monotonicity Constraint: To stabilize training and improve the representation of action advantages, a constraint is introduced. It penalizes cases where the estimated advantage (the difference between $R_Q(s, a)$ and $R_V(s)$) is negative, ensuring that the advantage function captures relative improvements over the state baseline.

Key Contributions

• SSRS Framework: A novel framework that integrates SSL into RL, enabling reward shaping on both non-zero and zero-reward transitions to improve sample efficiency in sparse-reward environments.
• Entropy Augmentation: The proposal of an entropy-based data perturbation method tailored for trajectory spaces, which outperforms traditional augmentation methods in non-visual observation tasks.
• Empirical Validation of SSL in RL: The paper empirically investigates the conditions under which SSL is effective in RL, demonstrating that trajectory spaces conform to the mild assumptions (smoothness and clustering) necessary for SSL methodologies.
• Monotonicity Constraint: The introduction of a constraint to regulate the relationship between state-value and state-action-value estimators, enhancing training stability.

Experimental Results
The SSRS framework was evaluated in sparse-reward Atari environments (e.g., Montezuma's Revenge, Seaquest, Venture) and robotic manipulation tasks (FetchReach).

• Performance vs. Baselines: SSRS consistently outperformed supervised-based approaches (RCP, SORS) and standard RL baselines. In highly sparse environments like Montezuma's Revenge, SSRS achieved peak scores up to twice those of supervised baselines.
• Augmentation Efficacy: The proposed entropy augmentation resulted in a 15.8% increase in the best score compared to other augmentation methods (e.g., cutout, smooth, translation).
• Ablation Studies: Experiments confirmed the importance of the monotonicity constraint, showing that SSRS with this constraint (SSRS-ST) achieved higher best scores than without it (SSRS-NST) in most tested environments.
• Hyperparameter Robustness: The study analyzed hyperparameters such as the update probability ($p_u$), confidence threshold ($\lambda$), and consistency coefficient ($\beta$). Results indicated that performance peaks when the confidence threshold is between 0.8 and 0.9, and the consistency coefficient has minimal impact on final results but affects training dynamics.

Significance and Claims
The paper claims that by treating reward shaping as a semi-supervised classification problem, agents can effectively learn from the majority of zero-reward transitions, which are typically discarded or underutilized. The authors argue that the trajectory space in RL satisfies the necessary assumptions for SSL (clustering and smoothness), allowing the model to infer reliable reward signals for zero-reward transitions. The introduction of entropy augmentation is highlighted as a critical innovation that adapts SSL perturbation techniques to the specific nature of RL observations, avoiding the pitfalls of applying image-centric augmentations to non-visual data. Ultimately, the work demonstrates that SSL can significantly enhance the efficacy of reward shaping, leading to higher agent scores in challenging, sparse-reward domains.