Designing a Reinforcement Learning Pipeline: Crafting the Reward Function and Simulator

Reinforcement Learning (RL) is a powerful method of artificial intelligence that more closely mirrors human decision making than other ML methods. With RL,  an agent  interacts with an environment iteratively, using each iteration to improve its decision-making abilities. 

This opens up the door to a number of possibilities that help AI tools (like Keebo) become better at optimizing cloud data warehouses:

• Responding to dynamic, evolving environments where Snowflake’s optimal state is always in flux 
• Navigating complex Snowflake environments with complex use cases and queries
• Exploring new optimization opportunities, rather than just exploiting previous knowledge

At the heart of every RL system lies the reward function and the reward simulator, which guide the agent toward achieving its goals. In this blog, we’ll break down how we design these two critical components to enable autonomous optimizations that reduce cost without hindering performance. 

1. Designing a well-defined reward function

The reward function is the “compass” that guides the RL agent. It tells the agent which behaviors are desirable and what are not. A well-designed reward function guides the agent toward the right decisions more efficiently, while a poorly designed reward function can steer the agent toward unintended or even harmful behaviours.

Here are some key principles we use at Keebo to design our reward function. 

a. Aligh rewards with the desired outcome

The reward function must reflect the ultimate goal of the task. For example:

Autonomous cloud data warehouse optimization should minimize costs while avoiding significant latency. The RL reward function should incentivize cost savings while also penalizing actions that cause query latency that falls outside designated parameters.

b. Avoid reward hacking

Reward hacking occurs when the agent finds loopholes to maximize rewards without actually achieving the desired outcome. For example:

In the cloud data warehouse scenario, the agent might reduce costs by aggressively downsizing resources. However, this could lead to severe latency spikes or system failures. To prevent this, the reward function must balance cost reduction with performance constraints.

c. Balance Immediate and Long-Term Rewards

RL agents often need to balance short-term gains with long-term success. For example:

In the cloud data warehouse context, the agent might save costs in the short term by reducing resources, but this could lead to long-term issues like degraded user experience or system instability. Use discounted rewards to encourage the agent to prioritize sustainable cost optimization. 

d. Keep It Simple

A complex reward function can make training unstable or slow. Start with a simple reward structure and refine it as needed. For example:

Instead of rewarding every tiny cost-saving action, reward only significant milestones, such as achieving a target cost reduction without violating latency thresholds.

Use case: Designing a reward function for cloud data warehouse optimization

• Positive rewards:
  • Reducing cloud infrastructure costs (e.g., scaling down unused resources)
  • Maintaining query latency within acceptable thresholds (e.g., under 500ms for 95% of queries)
• Negative rewards:
  • Exceeding latency thresholds (e.g., penalizing latency spikes above 1 second)
  • Over-provisioning resources unnecessarily (e.g., penalizing excessive resource allocation that leads to high costs)

2. Building a reward simulator

A reward simulator is a virtual environment where the RL agent can practice its decision-making and receive feedback on the quality of those decisions. It mimics the real-world scenario and calculates rewards based on the agent’s actions. A good simulator is crucial for training the agent efficiently and safely.

Here are the steps we take at Keebo to build our own reward simulator. 

a. Define the environment

The simulator must accurately represent the real-world environment where the agent will operate. This includes:

• States: The current configuration of the cloud data warehouse (e.g., number of nodes, CPU/memory usage, query latency).
• Actions: The choices available to the agent (e.g., scaling up/down resources, redistributing workloads).
• Transitions: How the environment changes in response to the agent’s actions (e.g., reducing nodes might increase latency or save costs).

b. Incorporate the reward function

The simulator uses the reward function to evaluate the agent’s actions. For example, in a cloud data warehouse, the simulator calculates rewards based on cost savings and latency metrics.  The RL algorithm may incur a positive reward by reducing costs by 10% while keeping latency under 500ms. On the other hand, an RL algorithm may incur a negative reward for causing latency to exceed 1 second, even if costs are reduced.

c. Ensure realism and scalability

The simulator should be realistic enough to prepare the agent for real-world challenges but also scalable to handle large amounts of training data. In our own simulators, we use historical data on query performance, resource usage, and costs to model the environment. Additionally, we make sure to simulate varying workloads (e.g., peak vs. off-peak hours) to ensure the agent can handle different scenarios.

d. Add randomness and noise

Real-world environments are frequently unpredictable. To account for this reality, we introduce randomness (e.g. sudden spikes in query load or temporary cloud service outages) to make the agent robust to uncertainty.

Use case: Building a simulator for cloud data warehouse optimization

• Environment: A virtual cloud data warehouse with simulated workloads, resource configurations, and query performance metrics.
• Rewards:
  • Points for reducing costs while maintaining latency within acceptable thresholds.
  • Penalties for causing latency spikes or over-provisioning resources.
• Randomness: Simulate unpredictable events like sudden workload spikes or temporary cloud service degradations.

Conclusion

Designing a well-defined reward function and building a reliable reward simulator are foundational steps in creating an effective reinforcement learning pipeline. For cloud data warehouse optimization specifically, the reward function needs to balance cost savings with performance expectations, while the simulator should model real-world cloud environments as accurately as possible.

By aligning rewards with desired outcomes, avoiding common pitfalls, and ensuring the simulator is both realistic and scalable, Keebo trains RL agents to optimize cloud costs effectively without compromising performance. 

To see Keebo’s RL in action, or to try out a free two-week trial, contact us here!