# Reward Augmentation in Reinforcement Learning for Testing Distributed Systems

arXiv: https://arxiv.org/abs/2409.02137 · HTML: https://mlsanity.com/paper/2409.02137/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL training of language and agent models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 9 of its 44 references that are in the map.
- First posted on arXiv 2 Sep 2024.

## Placed in

- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

## Builds on (most cited first, 9 of 9)

- [Reward Machines: Exploiting Reward Function Structure in Reinforcement Learning](https://mlsanity.com/paper/2010.03950/) (2020)
- [Reward-Free Exploration for Reinforcement Learning](https://mlsanity.com/paper/2002.02794/) (2020)
- [Compositional Reinforcement Learning from Logical Specifications](https://mlsanity.com/paper/2106.13906/) (2021)
- [Task-agnostic Exploration in Reinforcement Learning](https://mlsanity.com/paper/2006.09497/) (2020)
- [Transfer of Temporal Logic Formulas in Reinforcement Learning](https://mlsanity.com/paper/1909.04256/) (2019)
- [Temporal-Logic-Based Reward Shaping for Continuing Reinforcement Learning Tasks](https://mlsanity.com/paper/2007.01498/) (2020)
- [Deep Reinforcement Fuzzing](https://mlsanity.com/paper/1801.04589/) (2018)
- [MUZZ: Thread-aware Grey-box Fuzzing for Effective Bug Hunting in Multithreaded Programs](https://mlsanity.com/paper/2007.15943/) (2020)
- [Greybox Fuzzing of Distributed Systems](https://mlsanity.com/paper/2305.02601/) (2023)

## Abstract (arXiv)

Bugs in popular distributed protocol implementations have been the source of many downtimes in popular internet services. We describe a randomized testing approach for distributed protocol implementations based on reinforcement learning. Since the natural reward structure is very sparse, the key to successful exploration in reinforcement learning is reward augmentation. We show two different techniques that build on one another. First, we provide a decaying exploration bonus based on the discovery of new states -- the reward decays as the same state is visited multiple times. The exploration bonus captures the intuition from coverage-guided fuzzing of prioritizing new coverage points; in contrast to other schemes, we show that taking the maximum of the bonus and the Q-value leads to more effective exploration. Second, we provide waypoints to the algorithm as a sequence of predicates that capture interesting semantic scenarios. Waypoints exploit designer insight about the protocol and guide the exploration to ``interesting'' parts of the state space. Our reward structure ensures that new episodes can reliably get to deep interesting states even without execution caching. We have implemented our algorithm in Go. Our evaluation on three large benchmarks (RedisRaft, Etcd, and RSL) shows that our algorithm can significantly outperform baseline approaches in terms of coverage and bug finding.