Add RLHF guide and dummy demo with Keras/JAX #2117

TrailChai · 2025-06-02T16:06:51Z

This commit introduces a new example for Reinforcement Learning from Human Feedback (RLHF).

It includes:

`examples/rl/rlhf_dummy_demo.py`: A Python script demonstrating a simple RLHF loop with a dummy environment, a policy model, and a reward model, using Keras with the JAX backend.
`examples/rl/md/rlhf_dummy_demo.md`: A Markdown guide explaining the RLHF concept and the implementation details of the demo script.
`examples/rl/README.md`: A new README for the RL examples section, now including the RLHF demo.

Note: The Python demo script (`rlhf_dummy_demo.py`) currently experiences timeout issues during the training loop in the development environment, even with significantly reduced computational load. This is documented in the guide and README. The code serves as a structural example of implementing the RLHF components.

This commit introduces a new example for Reinforcement Learning from Human Feedback (RLHF). It includes: - \`examples/rl/rlhf_dummy_demo.py\`: A Python script demonstrating a simple RLHF loop with a dummy environment, a policy model, and a reward model, using Keras with the JAX backend. - \`examples/rl/md/rlhf_dummy_demo.md\`: A Markdown guide explaining the RLHF concept and the implementation details of the demo script. - \`examples/rl/README.md\`: A new README for the RL examples section, now including the RLHF demo. Note: The Python demo script (\`rlhf_dummy_demo.py\`) currently experiences timeout issues during the training loop in the development environment, even with significantly reduced computational load. This is documented in the guide and README. The code serves as a structural example of implementing the RLHF components.

divyashreepathihalli

Thanks Yasir! just keep the .py file for now. once it is approved we can generate the .ipynb files

examples/rl/README.md

examples/rl/md/rlhf_dummy_demo.md

divyashreepathihalli · 2025-06-03T18:20:02Z

examples/rl/rlhf_dummy_demo.py

+            policy_model_params["non_trainable"], 
+            state_input
+        )
+        actual_predictions_tensor = predictions_tuple[0] 


are we assuming batch size is 1?

Yes. The code as written assumes a batch size of 1 for all model inputs and gradient calculations

why is that?

This was a simplification for the demo's clarity and to manage complexity, especially since REINFORCE-style updates can be done on single trajectories. A more advanced setup would definitely use batching across multiple episodes or from a replay buffer for stability and efficiency.

Does that make sense in the context of this simplified demo?

divyashreepathihalli · 2025-06-03T18:21:56Z

examples/rl/rlhf_dummy_demo.py

+            episode_policy_losses.append(current_policy_loss)
+            policy_grads_step = policy_grads_dict_step["trainable"]
+            # Accumulate policy gradients
+            for i, grad in enumerate(policy_grads_step):


For potentially improving performance
policy_grads_accum = jax.tree_map(lambda acc, new: acc + new if new is not None else acc, policy_grads_accum, policy_grads_step)

Thanks, I have refactored policy gradients and reward gradients accumulations using jax.tree_map.

divyashreepathihalli · 2025-06-03T18:31:48Z

examples/rl/md/rlhf_dummy_demo.md

+        actual_predictions_tensor = predictions_tuple[0]
+        action_probs = actual_predictions_tensor[0]
+        log_prob = jnp.log(action_probs[action] + 1e-7)
+        return -log_prob * predicted_reward_value_stopped


If this predicted_reward_value is just R(s,a), then it's using the immediate predicted reward.
it's a very naive and generally ineffective way to train a policy. Thee log_prob should be multiplied by the cumulative discounted future reward (Return, G_t).

Commit c212593 addresses this.

…rewards. This commit refines the RLHF demo example (`examples/rl/rlhf_dummy_demo.py`) to use discounted cumulative actual rewards (G_t) for policy gradient calculations, aligning it with the REINFORCE algorithm. Changes include: - Added a `calculate_discounted_returns` helper function. - Modified the `rlhf_training_loop` to collect trajectories (states, actions, rewards) and compute G_t for each step at the end of an episode. - Updated the policy loss function to use these G_t values instead of immediate predicted rewards. - The reward model training logic remains focused on predicting immediate rewards based on simulated human feedback (environment reward in this demo). - Updated the corresponding RLHF guide (`examples/rl/md/rlhf_dummy_demo.md`) to explain these changes and provide updated code snippets. The timeout issues with the script in the development environment persist, but the code now better reflects a standard policy gradient approach.

divyashreepathihalli · 2025-06-03T23:58:59Z

the .md files are automatically generated. SO you might want to move the explanation content part to .py

Moving the explanation piece to .py from .md

TrailChai · 2025-06-10T12:59:00Z

I have deleted the .md files and added relevant documentation pieces to .py file

divyashreepathihalli · 2025-07-11T18:59:00Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a new example demonstrating a simplified Reinforcement Learning from Human Feedback (RLHF) loop using Keras with a JAX backend. The example includes a Python script with an embedded markdown guide.

My review focuses on improving clarity, correctness, and maintainability. Key findings include:

Inconsistent terminology (RLAIF vs. RLHF) that should be standardized.
A critical discrepancy between the documentation and the code regarding the policy update algorithm.
A bug in the logging logic that prevents any training progress from being displayed with the current settings.
A minor performance issue in a helper function and a formatting error in the documentation.

I've provided specific suggestions to address these points.

gemini-code-assist · 2025-07-11T19:00:57Z

examples/rl/rlhf_demo.py

+# Reinforcement Learning from AI Feedback(RLAIF) - Demo Guide
+This guide explains the concept of  Reinforcement Learning from AI Feedback (RLAIF) and walks through the components of the accompanying script `rlhf_demo.py`.


The guide's main title and introduction use "RLAIF" (Reinforcement Learning from AI Feedback), but the PR and the rest of the guide refer to "RLHF" (Reinforcement Learning from Human Feedback). This is inconsistent and confusing. Please use "RLHF" consistently. The metadata in lines 2 and 6 should also be updated to reflect this change.

Suggested change

# Reinforcement Learning from AI Feedback(RLAIF) - Demo Guide

This guide explains the concept of Reinforcement Learning from AI Feedback (RLAIF) and walks through the components of the accompanying script `rlhf_demo.py`.

# Reinforcement Learning from Human Feedback (RLHF) - Demo Guide

This guide explains the concept of Reinforcement Learning from Human Feedback (RLHF) and walks through the components of the accompanying script `rlhf_demo.py`.

gemini-code-assist · 2025-07-11T19:00:58Z

examples/rl/rlhf_demo.py

+            total_reward_loss_avg += mean_episode_reward_loss
+            loss_count_avg +=1
+
+        if (episode + 1) % 100 == 0 and loss_count_avg > 0:


The training progress is logged only when (episode + 1) % 100 == 0. However, the training loop is configured to run for only 10 episodes (both in the function default and the __main__ block call). As a result, no progress will be printed during the run. The logging interval should be adjusted to be useful for a short demo run.

Suggested change

if (episode + 1) % 100 == 0 and loss_count_avg > 0:

if (episode + 1) % 5 == 0 and loss_count_avg > 0:

gemini-code-assist · 2025-07-11T19:00:58Z

examples/rl/rlhf_demo.py

+4.  **Policy Model Training (REINFORCE-style):**
+    *   **At the end of each episode:**
+        *   The `calculate_discounted_returns` function is called with the list of `true_reward`s collected during the episode to compute the discounted cumulative reward (G_t) for each step.
+        *   These returns are typically normalized (subtract mean, divide by standard deviation) to stabilize training.
+        *   The code then iterates through each step `t` of the collected trajectory.
+        *   For each step, the `policy_loss_fn` is called. Its loss is calculated as `-log_prob(action_t) * G_t`. This means the update encourages actions that led to higher overall discounted future rewards.
+        *   Gradients for the policy model are computed for each step and accumulated across the entire episode.
+    *   **Gradient Application:** The accumulated policy gradients are averaged over the number of steps and applied to the policy model using its optimizer. This update rule aims to increase the probability of actions that lead to good long-term outcomes.


The documentation here describes a REINFORCE-style policy update where discounted returns (G_t) are calculated at the end of an episode from the true_rewards. However, the actual implementation in rlhf_training_loop performs an update at each step using the predicted_reward_value from the reward model as the signal for the policy update (see lines 192-208). The calculate_discounted_returns function is defined but never called.

This is a significant discrepancy that misleads users about how the algorithm is implemented. The documentation should be updated to accurately reflect the code's actor-critic-like implementation, where the reward model's output is used as the advantage signal for the policy update.

gemini-code-assist · 2025-07-11T19:00:58Z

examples/rl/rlhf_demo.py

+def calculate_discounted_returns(rewards, gamma=0.99):
+    returns = []
+    cumulative_return = 0
+    for r in reversed(rewards):
+        cumulative_return = r + gamma * cumulative_return
+        returns.insert(0, cumulative_return)
+    return jnp.array(returns)


The implementation of calculate_discounted_returns uses returns.insert(0, ...) inside a loop. Inserting at the beginning of a Python list is an O(n) operation, making the whole function O(n^2) where n is the number of rewards. While this function is not currently used, the PR description mentions timeout issues, so performance is a valid concern.

A more performant O(n) approach is to append to the list and then reverse it once at the end.

Suggested change

def calculate_discounted_returns(rewards, gamma=0.99):

returns = []

cumulative_return = 0

for r in reversed(rewards):

cumulative_return = r + gamma * cumulative_return

returns.insert(0, cumulative_return)

return jnp.array(returns)

def calculate_discounted_returns(rewards, gamma=0.99):

# A more performant O(n) approach is to append to the list and then reverse it.

returns = []

cumulative_return = 0

for r in reversed(rewards):

cumulative_return = r + gamma * cumulative_return

returns.append(cumulative_return)

return jnp.array(returns[::-1])

gemini-code-assist · 2025-07-11T19:00:58Z

examples/rl/rlhf_demo.py

+    *   **Gradient Application:** The accumulated policy gradients are averaged over the number of steps and applied to the policy model using its optimizer. This update rule aims to increase the probability of actions that lead to good long-term outcomes.
+5.  **Logging:** Average policy and reward losses for the episode are printed periodically.
+The core idea of RLHF is still present: we have a reward model that *could* be trained from human preferences. However, the policy update mechanism has shifted. Instead of using the reward model's output directly as the advantage signal for each step (as in the previous version of the script), the policy now learns from the actual discounted returns experienced in the episode, which is a more standard RL approach when actual rewards (or good proxies like `true_reward` here) are available for the full trajectory. In a full RLHF system, `episode_true_rewards` might themselves be replaced or augmented by the reward model's predictions if no dense "true" reward exists.
+8.  **Logging:** Periodically, average losses are printed.


The numbering in this list of key parts of the training loop is incorrect. It jumps from point 5 to point 8. This should be corrected to 6 for clarity.

Suggested change

8. **Logging:** Periodically, average losses are printed.

6. **Logging:** Periodically, average losses are printed.

github-actions bot assigned sachinprasadhs Jun 2, 2025

divyashreepathihalli reviewed Jun 3, 2025

View reviewed changes

TrailChai and others added 7 commits June 3, 2025 11:35

Delete examples/rl/README.md

af42d5f

Update and rename rlhf_dummy_demo.md to rlhf_demo.md

752829f

Update rlhf_demo.md

87d37f1

Update rlhf_dummy_demo.py

8e7711c

Delete examples/rl/README.md

2179fa4

Update rlhf_dummy_demo.py

dfea315

TrailChai requested a review from divyashreepathihalli June 3, 2025 22:11

TrailChai added 3 commits June 3, 2025 15:57

Update and rename rlhf_dummy_demo.md to rlhf_demo.md

637c615

Rename rlhf_dummy_demo.py to rlhf_demo.py

2e80261

Update rlhf_demo.md

4179b20

TrailChai added 2 commits June 6, 2025 23:06

Update rlhf_demo.py

0e97fa0

Moving the explanation piece to .py from .md

Delete examples/rl/md/rlhf_demo.md

30694ee

TrailChai added 3 commits June 18, 2025 19:38

Update rlhf_demo.py

622d342

Update rlhf_demo.py

da549a8

Update rlhf_demo.py

9c8094d

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

		# Reinforcement Learning from AI Feedback(RLAIF) - Demo Guide
		This guide explains the concept of Reinforcement Learning from AI Feedback (RLAIF) and walks through the components of the accompanying script `rlhf_demo.py`.

	if (episode + 1) % 100 == 0 and loss_count_avg > 0:
	if (episode + 1) % 5 == 0 and loss_count_avg > 0:

	8. Logging: Periodically, average losses are printed.
	6. Logging: Periodically, average losses are printed.

Add RLHF guide and dummy demo with Keras/JAX #2117

Are you sure you want to change the base?

Add RLHF guide and dummy demo with Keras/JAX #2117

Uh oh!

Conversation

TrailChai commented Jun 2, 2025

Uh oh!

divyashreepathihalli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyashreepathihalli Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

TrailChai Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

TrailChai Jun 4, 2025

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

TrailChai Jun 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

TrailChai Jun 4, 2025

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli commented Jun 3, 2025

Uh oh!

TrailChai commented Jun 10, 2025

Uh oh!

divyashreepathihalli commented Jul 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

TrailChai Jun 3, 2025 •

edited

Loading