fix: use jnp.exp(log_probs) instead of softmax(log_probs) in compute_entropy_from_logits

[kuishou68]

Summary

Fixes a bug in compute_entropy_from_logits in tunix/rl/ppo/ppo_helpers.py.

Closes #1386

Problem

Line 164 applies jax.nn.softmax(log_probs) to convert log-probabilities to probabilities. This is mathematically incorrect:

• softmax(log_softmax(x)) ≠ softmax(x)
• Only exp(log_softmax(x)) = softmax(x) (i.e., the true probability distribution)

Applying softmax to already log-normalized values produces a different distribution, leading to silently incorrect entropy values during PPO training.

Fix

Replace jax.nn.softmax(log_probs) with jnp.exp(log_probs) on line 164:

# Before (wrong):
probs = jax.nn.softmax(log_probs)

# After (correct):
probs = jnp.exp(log_probs)

Impact

The compute_entropy_from_logits function is called in ppo_learner.py to compute token_entropy. The bug causes incorrect entropy values, which can silently affect PPO training metrics and dynamics.