Observability tools for reinforcement learning

Introduction to reinforcement learning observability

The framework of mixed observable Markov decision processes (MOMDP)

Adapting reinforcement learning for partial observability

Monitoring vs. observability tools in machine learning

Essential features of a model observability solution

Real-time metric tracking: The tool should capture and display metrics from your model as it runs in real-time. In RL, for example, you’d monitor episode rewards, lengths, loss values, etc. Real-time tracking enables you to identify divergences or plateaus immediately, rather than discovering them after the fact.
Anomaly and drift detection: A robust observability platform will automatically check for unusual patterns, such as data drift (a change in input distribution) or performance anomalies. If your training reward suddenly drops or your model’s predictions shift compared to historical behavior, the tool should flag it. Early detection of anomalies helps prevent minor issues from turning into major failures.
Rich metadata and logging: Beyond just metrics, the solution should log metadata such as model versions, hyperparameters, dataset versions, and system information. It’s important to know which model (or which version of code) produced the results you’re seeing. Detailed logs (such as per-step actions in an RL episode, or intermediate outputs of a model) greatly aid in debugging complex behaviors.
Visualization and dashboards: Observability is made much easier by clear visualizations. A good solution provides dashboards where you can plot metrics over time, compare runs, and potentially visualize model internals (such as weights or feature activations). For RL, visual tools like viewing the agent’s policy improvement or seeing a heatmap of Q-values can be invaluable. Charts, histograms, and confusion matrices (for classification tasks) all help turn raw logs into insights.
Integration with MLOps workflows: The observability tool should integrate seamlessly with your existing machine learning operations. This means easy integration with training code (via simple APIs or callbacks), support for various environments (cloud, on-premise), and compatibility with pipeline orchestrators or CI/CD systems. It should output data in formats that can be consumed by other services or trigger alerts (e.g., send an email or Slack message if something goes wrong).
Scalability and collaboration: As a bonus, modern observability solutions are often cloud-based or have a web interface that teams can use. This means your whole team can collaborate, viewing the same dashboards and commenting on runs. The tool should handle numerous experiments or large datasets without bogging down. Scalability ensures that as your projects grow (with more models, more data, and more team members), the observability platform continues to perform optimally.

Impact of multimodal models and reinforcement learning on MLOps

Step-by-step tutorial: Using Weights & Biases for reinforcement learning observability

Step 1: Setting up Weights & Biases for observability

Install the W&B Python package. You can do this via pip. In a Jupyter notebook or script, run the following command to install W&B:

!pip install wandb

This will download and install the latest wandb package. (In a Jupyter environment, the ! at the start runs a shell command to install the package.)
Expected output:

Successfully installed wandb-0.15.8 etc.

The exact version number may vary, but you should see a success message after installation.
Login to W&B. After installation, you need to authorize W&B to use your account. You have a few options:
Easiest: In a notebook, call the login utility:

import wandb
wandb.login()

The first time, this will prompt you to enter your API key (it might display a link to get your key). Paste your API key and hit Enter. W&B will then save your credentials (typically in a local file ~/.netrc or environment variable) for future use.
Alternative: You can run wandb login from a terminal and paste your API key there, which achieves the same result.
If login is successful, W&B will confirm:

wandb: Logged in as <your_username>

Now you’re authenticated and ready to use W&B in your code.

Step 2: Monitoring model performance with Weights & Biases

Initialize the environment and Q-table.
Start a W&B run with wandb.init(), specifying the project and some hyperparameters to track.
Run a training loop for a number of episodes. In each episode, the agent will explore the environment and update its Q-table. We will log the episode reward to W&B at the end of each episode.
W&B will record those rewards, and we can watch the dashboard update with a reward curve as training progresses.
Finally, we’ll examine some output to confirm improvements.

!pip install gym           # Install OpenAI Gym for the environment
import gym
import numpy as np
import random
import wandb

# Initialize a new W&B run for logging
wandb.init(project="rl-observability-tools-tutorial", 
           config={ "episodes": 500, "learning_rate": 0.1, "discount_factor": 0.99, "epsilon_decay": 0.995 })

# Config parameters (retrieved from wandb.config for convenience)
episodes = wandb.config.episodes           # total training episodes
learning_rate = wandb.config.learning_rate # Q-learning learning rate
discount = wandb.config.discount_factor    # reward discount factor (gamma)
epsilon_decay = wandb.config.epsilon_decay # decay factor for exploration epsilon

# Set up the FrozenLake environment
env = gym.make("FrozenLake-v1", is_slippery=False)  # is_slippery=False for deterministic moves
num_states = env.observation_space.n    # number of states (16 for 4x4 FrozenLake)
num_actions = env.action_space.n        # number of possible actions (4: left, right, down, up)

# Initialize the Q-table to zeros
Q = np.zeros((num_states, num_actions))

# Exploration parameter (epsilon-greedy)
epsilon = 1.0  # start with full exploration
min_epsilon = 0.01

# Containers to track results for analysis
rewards_per_episode = []
success_per_episode = []  # 1 if goal reached in episode, 0 otherwise

for episode in range(1, episodes+1):
    state = env.reset()             # reset environment to start state
    total_reward = 0
    done = False

    # Run one episode
    while not done:
        # Choose action using epsilon-greedy policy
        if random.random() < epsilon:
            action = env.action_space.sample()   # explore: random action
        else:
            action = np.argmax(Q[state])         # exploit: best known action from Q-table

        # Take the action in the environment
        next_state, reward, done, info = env.step(action)

        # Update Q-value for the state-action pair using Q-learning update rule
        best_next_action = np.argmax(Q[next_state])
        td_target = reward + discount * Q[next_state][best_next_action]
        td_error = td_target - Q[state][action]
        Q[state][action] += learning_rate * td_error

        # Accumulate reward
        total_reward += reward
        state = next_state  # move to the next state

    # Episode finished
    rewards_per_episode.append(total_reward)
    success_per_episode.append(1 if total_reward > 0 else 0)  # success if reward 1 was obtained

    # Log the episode reward to W&B
    wandb.log({"episode": episode, "reward": total_reward})

    # Decay epsilon (less exploration as agent learns)
    if epsilon > min_epsilon:
        epsilon *= epsilon_decay

# Training loop finished
wandb.finish()  # end the W&B run and flush data

# Print final results
success_rate = sum(success_per_episode[-100:])  # successes in last 100 episodes
print(f"Run URL: {wandb.run.url}")
print(f"Success rate in last 100 episodes: {success_rate}%")

We installed and imported the necessary libraries. We use Gym for the environment, NumPy for numeric computations, and wandb for observability.
We call wandb.init(...) to start a W&B run. We gave the project a name "rl-observability-tools-tutorial" (you can choose any name; this will group your runs in the W&B web app). We also passed a config dictionary with some hyperparameters: number of episodes, learning rate, etc. W&B will save these as part of the run configuration, which is useful for record-keeping and comparison.
We configured the FrozenLake environment. Setting is_slippery=False makes the environment deterministic (the agent’s moves aren’t random), which helps the Q-learning converge faster for demonstration.
We initialized a Q-table (Q) with dimensions [num_states x num_actions] all to 0. This table will be learned over time.
We implemented an epsilon-greedy Q-learning loop:
For each episode, the agent resets to a start state and interacts with the environment until it reaches a terminal state (either falls in a hole or reaches the goal in FrozenLake).
At each step, with probability epsilon, the agent explores a random action; otherwise it exploits the current best action from the Q-table.
We take the action, observe the reward and next state, and update the Q-table using the Q-learning formula: [ Q[s, a] \leftarrow Q[s, a] + \alpha \Big( r + \gamma \max_a Q[s', a] - Q[s, a] \Big) ] where (\alpha) is the learning rate and (\gamma) is the discount factor.
We accumulate the reward for the episode and continue until the episode ends.
After each episode, we log the total_reward for that episode to W&B with wandb.log. We also track in lists whether the episode was a success (reached goal).
We decay epsilon after each episode to reduce exploration over time.
Once training is done, we call wandb.finish() to conclude the run. (This is generally optional as it will auto-finish when the script ends, but it’s good practice especially in notebooks to close the run explicitly.)
We then print out the W&B run URL and the success rate in the last 100 episodes.

Run URL: https://wandb.ai/your_wandb_username/rl-observability-tools-tutorial/runs/abc123def
Success rate in last 100 episodes: 100%

Step 3: Analyzing data and improving model observability

# Simple analysis of the logged data
print("First 10 episode rewards:", rewards_per_episode[:10])
print("Last 10 episode rewards:", rewards_per_episode[-10:])

average_reward = np.mean(rewards_per_episode)
success_rate = np.mean(success_per_episode) * 100
print(f"Overall average reward: {average_reward:.2f}")
print(f"Overall success rate: {success_rate:.2f}%")

First 10 episode rewards: [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
Last 10 episode rewards: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
Overall average reward: 0.56
Overall success rate: 56.00%

We can use the W&B UI to visualize the policy improvement over time. For instance, we could log the agent’s moving average reward and see when it crosses a desired threshold.
We could set up alerts in W&B (for example, an email or message if the reward drops below a certain value after a number of episodes, which might indicate a regression in learning).
We can analyze specific episodes. If an episode failed unexpectedly late in training, we could drill down into that episode. W&B allows logging custom data, such as the sequence of actions or states for an episode (using, for example, W&B Tables or uploading a video of gameplay in a more visual environment). In FrozenLake, we might not upload a video, but for environments like CartPole or Atari, recording a video of the agent every N episodes and logging it with Weights & Biases (wandb).Video would be very insightful.

If we notice that the learning is slow initially, we might try a higher learning rate or a different exploration schedule and run another experiment. With W&B, we can track the new run and then compare it side by side with other runs in the UI. For example, the reward curves of Run A and Run B could be overlaid to see which one learned faster.
We might use W&B’s Hyperparameter Sweep feature to systematically try combinations of parameters (learning_rate, epsilon_decay, etc.) to see which yields the best performance. The platform will automatically organize and visualize the results of the sweep.

Step 4: Alternative use cases and tips for maximizing effectiveness

Applying W&B to other ML tasks: Whether you’re training a computer vision model or a natural language model, the process of logging metrics and outputs is similar. For example, in a computer vision project, you might log the training and validation accuracy per epoch, and even sample predictions (images with the model’s labels vs. true labels) to a W&B dashboard. For an NLP model, you could log metrics like perplexity and sample-generated texts at different training stages. W&B’s observability features (charts, media panels, etc.) let you see how your model is performing beyond just final metrics – you can catch phenomena like overfitting by watching loss curves, or monitor a translation model’s outputs for quality over time.
Using W&B Weave for complex AI systems: W&B Weave is a suite of tools designed for advanced model observability, particularly for systems such as large language model agents or multimodal pipelines. If your project involves an agentic system (e.g., an AI agent that takes actions based on LLM outputs, or an RL agent interacting with an environment using language), Weave provides observability features tailored to that specific use case. It can trace interactions, log prompts, and responses for LLMs, and monitor long-running agent behaviors. Weave’s Monitors feature allows you to set up continuous monitoring in production – for example, keeping an eye on an RL-driven service and detecting anomalies in behavior or performance drift. Weave also includes Guardrails to catch harmful or unexpected outputs (more relevant in NLP, but conceptually, guardrails in an RL setting could mean ensuring the agent doesn’t take actions outside a safe set). All these help ensure that even as your AI systems get more complex, you maintain visibility into their workings.
Leveraging the model registry and artifacts: Weights & Biases includes a Model Registry (under the W&B Models section) and Artifacts system. These are incredibly useful for observability in a larger MLOps pipeline. After you’ve trained an RL model (or any model), you can use the registry to version and store your model. Each model version can be linked to the W&B run that produced it, which means you have a complete record of metrics, hyperparameters, and even dataset versions (if logged as artifacts) for that model. This traceability is gold for observability – when a model is deployed, you can always trace back to see how it was trained and how it performed then. Artifacts, on the other hand, enable you to track data and model files throughout training workflows. For example, you might log the dataset you used for training as an artifact and the final model weights as another artifact. This makes your experiments reproducible and debuggable: if something goes wrong with Model X in production, you can retrieve the exact training data and weights from W&B and investigate.
Hyperparameter tuning and experiments at scale: As mentioned, W&B Sweeps can automate the running of multiple experiments with different hyperparameters. This is an extension of observability – you’re not just observing one training run, but an entire family of runs to see how variations affect performance. The W&B interface will aggregate results of sweeps, letting you drill down into the best runs and compare them. Embracing such experimentation is key to improving model performance, and an observability platform makes it manageable by collecting all the metrics in one place. You can tag or group runs (for example, tag all runs that use Algorithm A vs Algorithm B) and then use the UI to filter and analyze groups.
Setting alerts and automations: W&B allows you to set alerts (under Settings or via the API) for certain conditions. For instance, you could set an alert if the reward in your RL training run hasn’t improved for a certain number of steps (indicating potential stagnation), or if the loss goes NaN (indicating a crash or divergence). In a production ML setting, you could set alerts on a drop in accuracy or a spike in error rate. These alerts can be sent to your email or messaging apps, acting as an early warning system. Additionally, W&B’s Automations can trigger actions – for example, automatically kick off a new training run when an upstream dataset artifact is updated, or retrain a model if performance dips. This type of automation, combined with observability, enables a robust MLOps pipeline where monitoring results can directly lead to remedial actions.
Collaboration and reporting: One often underappreciated aspect of using a tool like W&B is how easy it becomes to share results and insights. You can create Reports in W&B, which are like interactive articles or dashboards summarizing your findings (much like this tutorial itself). These can be shared within your team or publicly. For example, after running experiments, you could make a report comparing different algorithms on your task, with charts and tables all pulled from the W&B runs. This makes collaboration more effective – instead of everyone running code to see results, the results are accessible and nicely formatted for discussion. For observability, this means important discoveries (like “Algorithm A is more stable but slower than Algorithm B” or “Model began to overfit after epoch 5”) are clearly communicated.

Sources

Hai Nguyen et al., “Hierarchical Reinforcement Learning under Mixed Observability” – arXiv:2204.00898 (2022). (arxiv.org)
Stephan Weigand et al., “Reinforcement Learning using Guided Observability (PO-GRL)” – arXiv:2104.10986 (2021). (arxiv.org)
Chip Huyen, Monitoring Machine Learning: Interview with Oren Razon, ML in Production blog (2020). (mlinproduction.com)
Weights & Biases – W&B Weave Product Page, wandb.ai (2024). (wandb.ai)
Weights & Biases – LLM Observability with W&B Weave, Fully Connected blog (2023). (wandb.ai)
Stable Baselines3 Documentation – W&B Integration Callback. (stable-baselines3.readthedocs.io)