This is the second article of three. In the first article, “Connecting Java Reinforcement Learning to Python Gymnasium“, we covered the integration between Java and Gymnasium and established the foundation for testing the algorithm in simple environments such as CartPole and MountainCar before moving to CARLA. This article covers the implementation of DQN to Rainbow DQN (using Deep Java Library1); this algorithm will be the building core for the autonomous car in the third article.
Table of Contents
- Memory Model
- Calculus
- DQN – Deep-Q Network
- Core Architecture: Online and Target Networks
- Training Loop, Exploration, and Replay Buffer
- Q-Value and Target Network Update Schedule
- DQN Training and Inference Flow Diagrams
- Abstract Agent Training Method Implementation
- Cleanup, Debugging, and Online Training Override
- DQN Formula and Q-Table vs Neural Network Approach
- Stability via Target Network and Replay Buffer Sampling
- Gather Operation and Index Array Explanation
- Agent Train Method Implementation Details
- Target Network Forward and Cloning
- Architecture Definition
- Model Loading, Weight Initialization, and Backward Loss
- Execution and Configuration Setup
- DDQN – Double DQN
- Dueling DQN
- PER – Prioritized Experience Replay Buffer
- N-Step Learning
- Noisy Networks
- Noise Distributions and the Reparameterization Trick
- Custom Implementation in DJL
- Factorized Noise: Efficient Sampling via Outer Product
- Verifying Noise Samples Are Not Trainable
- Step-by-Step Sample Generation Process
- Efficiency Comparison: Factorized vs. Direct Sampling
- Computational Efficiency: Factorized vs. Independent Sampling
- Initialization Strategy for Noisy Networks
- Noisy Initialization Implementation
- Building the Noisy Network and Training Loop
- Shannon Entropy and Kullback-Leibler Intuition
- Distributional DQN (C51)
- Original Distributions
- Bellman Projection
- Bellman Projection Java Implementation
- Impact of Atom Count on Distribution Resolution
- Network Implementation in DJL
- Categorical C51 Loss Implementation
- Training Loop
- Step-by-Step Value Transformation Explanation
- Atom Projection: Shift and Clip Each Atom
- Mass Redistribution via Linear Interpolation
- KL Divergence and Cross-Entropy in C51
- Distributional RL vs. Standard DQN
- Architecture Comparison: DQN, Dueling DQN, and C51
- Rainbow DQN
Rainbow DQN combines six improvements over the original DQN rather than introducing a new algorithm: Prioritized Experience Replay (PER), Dueling DQN, Double DQN, N-step returns, NoisyNets, and C51 (Categorical DQN). Each improvement addresses a specific limitation of the original algorithm, and together they produce a significantly more capable agent.
The implementation follows the same project structure from the first article, java-rl-dqn-to-rainbow, where I add each component incrementally, keeping the algorithm testable at every stage. I used the rainbow-is-all-you-need project as my reference implementation.
The main environment that we’ll use to train our models will be Lunar Lander2, because it’s fast to train and easy to check if it’s learned. The gif below is from our trained AgentDQN model:
Let’s start with the memory model and main methods used to build this implementation.
Memory Model
DJL implements scope-based memory management, similar to the region-based approach3, where memory managers define separate scopes for systems or blocks. For example, in the code below, we have the parent memory manager created inside a broader block, and inside each inner block, we allocate new memory managers using the method subMrg(). The colors at the end of each block show when the sub manager deallocates. After all the blocks run, all the off-heap memory deallocates. If, for example, you deallocate parent directly, you automatically deallocate all the memory of sub managers.
The subMgr() Method Implementation
// class DJLMemoryManagement
public static NDManager subMgr(NDManager manager, String name) {
var sub = manager.newSubManager();
sub.setName(name + “-” + sub.getName());
return sub;
}
The Problem with Nested Function Calls
Another problem with memory management that we can have when working with native objects (allocated off-heap memory), is invoking nested function calls. In the example below we allocate 2 new objects, one for the mul(…) and another for the sum(…):
dist.mul(sup).sum(LAST_AXIS_ARR);
The scoped() Solution
How do we solve this? Using the following method:
// class DJLMemoryManagement
public static NDArray scoped(final Function<NDArray[], NDArray> block,
final NDArray... arrays) {
if (arrays.length == 0) throw new IllegalArgumentException(“arrays must contain elements!”);
@Cleanup var sub = subMgr(arrays[0], “scoped-arrays”);
sub.tempAttachAll(arrays);
var result = sub.ret(block.apply(arrays));
for (var element : arrays) {
if (element == result) {
throw new IllegalStateException(“scoped block returned any input NDArray itself”);
}
}
return result;
}
We create a temporary sub memory manager, and temporarily attach (by borrowing) all elements to it. After that, every operation that creates new elements will use this sub NDManager as its main manager. We must call sub.ret(…) because we want the output of the lambda block.apply(…) to survive this scope and remain usable outside this method. When sub is closed, all borrowed NDArrays are returned to their original NDManager owners.
Applying the Solution Locally
Applying the solution locally avoids memory leaks, as described below.
// class CategoricalQNetworkMLP
return scoped(array -> {
var dist = array[0];
var sup = array[1];
return dist.mul(sup).sum(LAST_AXIS_ARR);
}, distribution, support);
Strategic System-Level Memory Management
You could let the parent memory managers hold these temporary objects, and let the parent deallocate the memory (using try-with-resources or @Cleanup), and consider the purpose of each system. For example, System B takes care of its own memory, the same with System C, where they are both children of System A. If a system performs excessive allocations, consider introducing smaller sub-systems to avoid Out of Memory Error, limiting its lifetime.
Bad vs Good Practices
Check the example below of bad vs good practices:
Memory Model of the Project
Using the DJL memory management4 as base, we built an adaptation to manage the memory internally, below you can see the diagram of the tree structure of the NDManagers (parent and its children):
Note that DeepQNetworkMLP and IReplayBuffer have a longer lifecycle, but internally, they are implemented to release temporary objects after each execution, using the practices mentioned before. Below you can see a safe training/inference:
// class DJLMemoryManagement
public static NDList safeForwardSingle(NDManager manager,
Block block,
ParameterStore parameterStore,
NDList inputs,
boolean training) {
if (inputs.size() != 1) throw new IllegalArgumentException("The network just accepts one head input!");
@Cleanup var sub = subMgr(manager, "safe-forward-single");
inputs.getFirst().tempAttach(sub);
var output = block.forward(parameterStore, inputs, training);
output.attach(manager);
return output;
}
Now focusing on the loop, you see that we have a sub manager per episode, and a sub manager for each frame too, with this pattern, if an inner operation fails to explicitly release temporary memory, the NDManager above in the hierarchy can release it after exiting the loop (or any scope that it’s in):
AbstractAgent loop.Now that we concluded the memory management model, let’s understand and build DQN and its derivatives, step by step.
Calculus
The article below gives you an introduction to calculus and a linear regression. This serves as a base to understand how derivation works and help you understand the formulas and code that I’ll present later in this second article about reinforcement learning using Java:
- Linear Regression Using the Ordinary Least Squares Method
- In this article we’ll build the Ordinary Least Squares Method step by step, explaining why we use slope-intercept form and derivatives, then we’ll execute Python code to visualize how this function behaves.
- The videos below can help you understand and remember the RMSD, Chain Rule and Multi-Variable Calculus:
DQN – Deep-Q Network
To solve problems in complex spaces, like a game, we use neural networks as function approximators; DQN using its neural network can work with ample types of problems.
Core Architecture: Online and Target Networks
We first create two neural networks, an online-net and a target-net. The target-net starts as a clone of the online-net, but we freeze its weights (to serve as a stable bootstrap) and do not modify them during backpropagation. This approach gives the agent a fixed baseline for comparing its values during training. If we trained the target net as well, both networks would have shifting values, which would destabilize training and make learning much harder, or even impossible.
Training Loop, Exploration, and Replay Buffer
The agent learns within an environment, where it takes actions. At the beginning, it acts more randomly to encourage exploration (sacrificing immediate reward for potentially better future rewards). Later, the agent increasingly bases its actions on what it has learned, shifting toward exploitation (now that it has found a good reward, it can repeat that successful behavior more often). After each action is taken in the environment, the state, action, reward, next state and done flag are saved on the Replay Buffer. The Replay Buffer is responsible for allowing the networks to decorrelate training (if sampling is done properly) and mitigate catastrophic forgetting (where the network unlearns things); when we have enough data, we train our online-net using TD-Error (TD is Temporal Difference), where we get the difference between the actual Q-Value and TD-Target (explained in detail later).
Q-Value and Target Network Update Schedule
The Q-Value is a representation of the expected discounted accumulated reward when you are in some state s, taking action a.
When do we update the target-net? After C time steps, where C can be the number of frames or episodes (episodes length can vary, which is imprecise to measure between runs).
An episode is the first step after the initialization or reset until you receive a done for the environment.
DQN Training and Inference Flow Diagrams
The mermaid diagram below shows the process:
During the inference process you don’t explore anymore, just exploit the environment, getting the best actions the agent knows:
Abstract Agent Training Method Implementation
The code below is the implementation of the explained flow above, that most agents based on DQN (and its derivatives) use:
// class AbstractAgent
@Override
public void train(int batchSize,
long framesLimit,
int framesSkip,
IReplayBuffer replayBuffer,
Loss lossFunc) {
// omitted ...
plotTrackers.explainTrackersMessage();
this.test = false;
this.episodes = 0;
@Cleanup var pbar = ProgressBar.builder()
.setTaskName("training " + getClass().getSimpleName())
.setInitialMax(framesLimit)
.setMaxRenderedLength(200)
.build();
int frames = 0;
do {
@Cleanup var parentPerEpisode = subMgr(parent, "parent-episode");
var stateAndInfoMap = env.reset(parentPerEpisode);
@Cleanup var state = stateAndInfoMap.state();
var episodeRewards = new ArrayList<Double>();
float episodeLossSum = 0f;
int episodeSteps = 0;
do {
@Cleanup var sub = subMgr(parentPerEpisode, "sub-frame");
@Cleanup var action = selectAction(state);
var stepResult = env.step(action, sub);
var reward = stepResult.reward();
var nextState = stepResult.state();
var done = stepResult.done();
var info = stepResult.info();
if (!info.isEmpty()) lastInfo = info;
var exp = new Experience(state.duplicate(), action.duplicate(), reward,
nextState.duplicate(), done);
replayBuffer.store(exp);
float lossItem = trainOnline(batchSize, replayBuffer, lossFunc, sub);
if (trained(lossItem)) {
episodeLossSum += lossItem;
++episodeSteps;
}
frames += framesSkip;
updateTargetNetworkAtN(frames);
episodeRewards.add(reward);
if (done) {
++episodes;
float avgLoss = episodeSteps > 0 ? episodeLossSum / episodeSteps : 0;
plotTrackers.add(epsilon, episodeRewards, avgLoss);
release(state, nextState);
break;
}
state = transfer(parentPerEpisode, state, nextState);
pbar.stepTo(frames);
plotTrackers.updateTrackersMessage(pbar, frames, replayBuffer.size(), parent, parentPerEpisode);
templateExtraProcessing(frames, framesLimit);
epsilon = reduceEpsilon(epsilon);
} while (frames < framesLimit);
} while ((frames < framesLimit));
plotTrackers.updateTrackersMessage(pbar, frames, replayBuffer.size(), parent);
if (debugMemoryLeak) {
this.managerNode = getDebugDump(parent).orElse(null);
}
if (getBoolProp("agent.showAllMetrics")) {
plotTrackers.showAllMetrics();
}
}
Cleanup, Debugging, and Online Training Override
The @Cleanup is used to reduce the verbosity of try-with-resources. You may note the method getDebugDump(); this method checks for native resources allocated in PyTorch (or TensorFlow, MXNet) that remain undeallocated.
The trainOnline will be overridden by the child classes later:
// class AbstractAgent
protected abstract float trainOnline(int batchSize,
IReplayBuffer replayBuffer,
Loss lossFunc,
NDManager sub);
Then we have the methods to update the target-net every C frames and the method that explores or exploits the environment respectively:
// class AbstractAgent
protected void updateTargetNetworkAtN(int frames) {
if (frames % updateQTargetAtTimeN == 0) {
copy(onlineNet.getBlock(), targetNet.getBlock());
freeze(targetNet.getBlock());
}
}
@Override
public ActionSpaceType.ActionResult selectAction(NDArray state) {
if (!test) {
var rand = ThreadLocalRandom.current().nextFloat();
if (rand < epsilon) {
return env.actionSpaceSample();
}
}
@Cleanup var oneBatchState = state.expandDims(0);
long action = onlineNet.forwardLong(oneBatchState,
qValue -> qValue.stopGradient().argMax(AXIS_1));
return actionSpaceType.get(action);
}
DQN Formula and Q-Table vs Neural Network Approach
Let’s proceed to the base DQN formula:
DQN5 is an extension of Q-Learning6: instead of a Q-table, which becomes impractical when the state space is too large or continuous, a neural network approximates the action-value function Q∗(s,a). In the original paper, the network was a CNN (Convolutional Neural Network) trained directly on raw game pixels, with the input preprocessed into a stack of 4 grayscale frames of shape 84 x 84, providing temporal information.
Stability via Target Network and Replay Buffer Sampling
Two networks are used: the online network, updated at every training step via backpropagation, and the target network, a frozen copy updated only periodically by copying the online network’s weights, at C steps. Without this separation, training becomes unstable; the network would chase a target that moves at the same rate as its own updates, which creates the “dog chasing its tail” problem.
Past experiences (s,a,r,s’,done) are stored in a replay buffer and randomly sampled during training, reducing the correlation between consecutive experiences. For each sample, the target y is computed using the target network, and the online network is trained to minimize:
Notice that the gradient passes only through the online network. The TD-Target remains a constant during this update because separate, frozen parameters fix the target network: θ⁻ ≠ θ
Gather Operation and Index Array Explanation
The index array below, named actions, is used in the gather(…) method; it needs to have the same number of dimensions as the array you’re gathering from. Check below the index array values:
ND: (6, 1) gpu(0) int64
[[ 0],
[ 1],
[ 2],
[ 3],
[ 4],
[ 1],
]
Now the image illustrates how each position is accessed and referenced in the result variable:
Agent Train Method Implementation Details
The implementation of the trainOnline method checks if we have enough samples in replayBuffer; if true, we apply the formulas above. We used the comments both to implement the algorithms and to aid in teaching.
// class AgentDQN
@Override
protected float trainOnline(int batchSize,
IReplayBuffer replayBuffer,
Loss lossFunc,
NDManager sub) {
if (!replayBuffer.enough(batchSize)) return Float.NaN;
@Cleanup var samples = replayBuffer.sample(batchSize);
@Cleanup var targetQValue = targetNet.forward(samples.nextStates(), nextQValue -> {
// max Q(s’, a’)
var maxNextQValue = nextQValue.max(AXIS_1_ARR, KEEP_DIMS);
// gamma * max Q(s’, a’)
var discountNextQValue = maxNextQValue.mul(gamma);
// (1 - done)
var mask = samples.dones().neg().add(1);
// r + gamma * max Q(s’, a’) * (1 - done)
return samples.rewards()
.add(discountNextQValue.mul(mask))
.stopGradient();
});
float lossItem = backwardLoss(sub, lossFunc, targetQValue, () -> {
var states = samples.states();
var actions = samples.actions();
// y_hat = q_online(s, a)
return onlineNet.forward(states, qValue -> qValue.gather(actions, AXIS_1));
});
trainStep(onlineNet.getBlock(), optimizer);
return lossItem;
}
Target Network Forward and Cloning
The targetNet.forward(…) uses the safeForwardSingle(…) so any NDArray created by the targetNet will be safely deallocated inside the block.
The clone() is responsible for creating the target-network.
Architecture Definition
The neural network is a SequentialBlock that extends the class AbstractBlock; this architecture allows you to compose blocks within blocks, allowing flexibility. As the Lunar Lander problem is simple, we define an architecture with 3 layers only:
// class DeepQNetworkMLP
// omitted code ...
@SneakyThrows
public DeepQNetworkMLP(int observations,
int actions,
Path modelPath,
String prefix,
NDManager parent) {
this.observations = observations;
this.actions = actions;
this.subManager = subMgr(parent, getClass());
this.model = newModel(getClass(), subManager.getDevice());
this.net = new SequentialBlock();
// neural network architecture
net.add(linear(128))
.add(Activation::relu)
.add(linear(128))
.add(Activation::relu)
.add(linear(actions));
model.setBlock(net);
this.parameterStore = new ParameterStore(subManager, false);
if (modelPath != null) {
log.info(“Loading model: {}, {}”, modelPath, prefix);
model.load(modelPath.normalize(), prefix);
this.training = false;
} else {
net.initialize(subManager, DataType.FLOAT32, new Shape(1, observations));
initGradients(model.getBlock());
this.training = true;
}
}
// omitted code ...
@Override
public NDList forward(NDList input) {
return safeForwardSingle(subManager, net, parameterStore, input, training);
}
@Override
public NDArray forward(NDArray input) {
return forward(new NDList(input)).singletonOrThrow();
}
// omitted code ...
@Override
public IDeepQNetwork clone() {
var cloned = new DeepQNetworkMLP(observations, actions, subManager.getParentManager());
setName(cloned.subManager, "clone");
copy(model.getBlock(), cloned.model.getBlock());
return cloned;
}
// omitted code ...
Model Loading, Weight Initialization, and Backward Loss
The model serves as a base for the neural network architecture, and contains methods to load trained models or initialize the weights.
The method backwardLoss(…) and its derivatives are used to backpropagate and train the onlineNet. What allows the gradient calculation inside this method is the gradCol of type GradientCollector:
// class DJLLoss
public static float backwardLoss(final NDManager manager,
final Loss lossFunc,
final NDArray yTarget,
final Supplier<NDArray> yPredBlock) {
@Cleanup var sub = subMgr(manager, “scoped-back”);
@Cleanup var gradCol = gradient();
var yPred = yPredBlock.get();
var lossesVal = evaluate(lossFunc, sub, gradCol, yTarget, yPred);
return lossesVal.stopGradient().mean().getFloat();
}
Note the @Cleanup applied to the sub manager. Without closing this NDManager at the end of the scope, the gradients and other resources associated with backward computation can accumulate across training steps, making the training diverge eventually. In this implementation, the problem is contained by using smaller memory scopes, such as a dedicated NDManager for each frame.
The trainStep(…) method updates only the parameters whose NDArray has gradient computation enabled in the network. This also makes this method suitable for transfer learning7, where specific layers or parameters can be frozen while the remaining parameters are trained:
// class DJLOptimizer
public static boolean trainStep(Block block, Optimizer optimizer) {
boolean atLeastOneGradient = false;
for (var param : block.getParameters()) {
if (!param.getValue().requiresGradient()) continue;
atLeastOneGradient = true;
var weight = param.getValue().getArray();
@Cleanup var grad = weight.getGradient();
optimizer.update(param.getKey(), weight, grad);
}
if (!atLeastOneGradient) log.warn(“No training was made!”);
return atLeastOneGradient;
}
ExperienceReplayBuffer is the first IReplayBuffer implementation, as stated before, it’s used to decorrelate data between training, as we make a random sampling in the buffer:
// omitted code ...
public class ExperienceReplayBuffer implements IReplayBuffer {
private final Experience[] experiences;
private final NDManager subManager;
private final ExperienceSampler sampler;
private final int capacity;
private int size;
private int pos;
// omitted code ...
@Override
public void store(Experience exp) {
exp.state().attach(subManager);
exp.nextState().attach(subManager);
if (size < capacity) ++size;
var oldExp = experiences[pos];
experiences[pos] = exp;
setName(exp.state(), “buffer-state”);
setName(exp.nextState(), “buffer-next-state”);
if (oldExp != null) oldExp.close();
pos = (pos + 1) % experiences.length;
}
@Override
public boolean enough(int batchSize) {
return size >= batchSize;
}
@Override
public VecExperience sample(int batchSize) {
if (!enough(batchSize)) return null;
var sub = subMgr(subManager, “buffer-sample”);
var batch = sampler.sample(experiences, batchSize, size, false);
var states = newAttachedList(sub, batch, exp -> exp.state().duplicate());
var actions = djlMapToLong(sub, batch, exp -> exp.actionAs(Long.class));
var rewards = djlMapToFloat32(sub, batch, Experience::reward);
var nextStates = newAttachedList(sub, batch, exp -> exp.nextState().duplicate());
var dones = djlMapToFloat32(sub, batch, exp -> exp.done() ? 1 : 0);
return new VecExperience(
sub,
states,
actions,
rewards,
nextStates,
dones
);
}
@Override
public int capacity() {
return capacity;
}
@Override
public int size() {
return size;
}
int pos() {
return pos;
}
public record VecExperience(NDManager sub,
NDArray states,
NDArray actions,
NDArray rewards,
NDArray nextStates,
NDArray dones) implements IVecExperience {
@Override
public void close() {
sub.close();
}
}
@Override
public boolean isOpen() {
return subManager.isOpen();
}
@Override
public void close() {
release(experiences);
subManager.close();
}
}
The code below makes sure that the duplicated state and nextState are now owned by the subManager of ExperienceReplayBuffer, because we want these two objects to survive while we are training:
exp.state().attach(subManager);
exp.nextState().attach(subManager);
The VecExperience receives in its constructor the temporary sub manager (with the name pattern buffer-sample in it); this allows us to deallocate every object associated with the sample that itself is temporary too, because it’s a copy of the elements presented in the buffer.
Execution and Configuration Setup
Then we finish DQN part with the actual execution of the agent, where you can train and then run the algorithm:
// omitted code ...
public class AgentDQNMain {
static void main() {
run();
}
static void main(String[] args) {
var opts = parseArgs(args, AgentDQNMain.class.getSimpleName());
run(opts);
}
public static Optional<DJLMemoryManagement.ManagerNode> run() {
return run(RLRunOptions.defaults());
}
public static Optional<DJLMemoryManagement.ManagerNode> run(RLRunOptions opts) {
var builder = RLConfig.builder()
.envName("LunarLander-v3")
.runnerClass(AgentDQNMain.class.getSimpleName())
.algorithmName("dqn")
.observations(8)
.actions(4)
.learningRate(0.0005f)
.maxEpsilon(1.0f)
.minEpsilon(0.01f)
.discountFactor(0.99f)
.updateQTargetAtTimeN(1000)
.batchSize(128)
.framesLimit(getIntProp("agent.framesLimit", "300000"))
.bufferCapacity(getIntProp("agent.bufferCapacity", "30000"))
.saveModel(getBoolProp("agent.saveModel", "true"))
.debugMemoryLeak(getBoolProp("agent.debugMemoryLeak", "true"))
.renderRun(getBoolProp("agent.renderRun", "true"))
.runMaxTries(getIntProp("agent.maxTries", "1"));
if (opts.loadModelPrefix() != null) {
builder = builder.loadModelPrefix(opts.loadModelPrefix());
}
var config = builder.build();
return RLRunner.run(config, (env, optimizer, plotTrackers, parent) ->
buildDQN(config, env, optimizer, plotTrackers, parent));
}
private static IAgent buildDQN(RLConfig config,
IEnv env,
Optimizer optimizer,
PlotTrackers plotTrackers,
NDManager parent) {
boolean loadModel = config.loadModelPrefix() != null;
Supplier<IDeepQNetwork> networkFactory = loadModel?
() -> new DeepQNetworkMLP(
config.observations(),
config.actions(),
config.path(),
config.loadModelPrefix(),
parent)
: () -> new DeepQNetworkMLP(
config.observations(),
config.actions(),
parent);
return new AgentDQN(
config.maxEpsilon(),
config.updateQTargetAtTimeN(),
config.minEpsilon(),
config.epsilonDecay(),
config.discountFactor(), // gamma
env,
optimizer,
parent,
networkFactory,
plotTrackers,
config.debugMemoryLeak()
);
}
}
The pattern Agent*Main is repeated for the other agents, with some particular differences, like exchanging the IReplayBuffer or IDeepQNetwork implementation, or even some specific tuning. You can check the repository.
DDQN – Double DQN
Traditional DQN suffers from an overestimation bias. Because the target always maximizes the Q-value, this bias creeps into your estimates. This gets worse when there’s noise, for instance, when an action randomly receives a very high reward (e.g. investing in a company that gains money by faking or when you won a game because of a bug); this wrong estimation propagates through the network. The solution to reduce the probability of this happening is to use the Double-DQN8, shown below:
Note that we make the selection with the Q-online network, while we perform the evaluation with the Q-target network. A simple analogy would be that when you see a problem, and your colleague sees a different solution, you estimate his solution based on your experience; it’s different when you always choose the same solution and evaluate this yourself. Just because you know some tool well doesn’t mean that it’s good enough for every situation.
Check the implementation below:
// class AgentDDQN
@Override
protected float trainOnline(int batchSize, IReplayBuffer replayBuffer, Loss lossFunc, NDManager sub) {
if (!replayBuffer.enough(batchSize)) return Float.NaN;
@Cleanup var samples = replayBuffer.sample(batchSize);
// a* = arg max q_online(s’, a’)
var nextStates = samples.nextStates();
@Cleanup var action = onlineNet.forward(nextStates,
onlineQNextValue -> onlineQNextValue.stopGradient().argMax(AXIS_1)
.reshape(N_BATCH, 1));
@Cleanup var targetQValue = targetNet.forward(nextStates, nextQValues -> {
// q_target(s’, a*) - DDQN
var qNextValue = nextQValues.gather(action, AXIS_1);
// gamma * q_target(s’, a*)
var discountNextQValue = qNextValue.mul(gamma);
// (1 - done)
var mask = samples.dones().neg().add(1);
// y = r + gamma * q_target(s’, arg max q_online(s’, a’)) * (1 - done)
return samples.rewards()
.add(discountNextQValue.mul(mask))
.stopGradient();
});
float lossItem = backwardLoss(sub, lossFunc, targetQValue, () -> {
var states = samples.states();
var actions = samples.actions();
// y_hat = q_online(s, a)
return onlineNet.forward(states, qValue -> qValue.gather(actions, AXIS_1));
});
trainStep(onlineNet.getBlock(), optimizer);
return lossItem;
}
Dueling DQN
Another problem we have with the DQN is that we can’t know the expected return of the state, represented by V(s) (value of state s); the Dueling DQN9 solves this with the formulas:
The paper states that the mean version is more common. We show the difference between them below:
When we use the mean, we can have values that are greater than or less than the mean. With max subtraction, the best action gets zero advantage, and other actions would get negative advantage values for suboptimal actions. When we encounter a situation where taking an action won’t make a difference, we would have the scenario where all advantage values are zero, and only the V(s) contributes to the Q-value.
Another great attribute of the Dueling-DQN formula is that it is identifiable. I’ll present the problem Dueling-DQN solves with numbers, check the example of an unidentifiable decomposition below:
You can see that without a constraint, we cannot know how much each V(s) and A(s, a) contributed to Q(s, a), but by subtracting the average (or max), the function becomes identifiable:
The Dueling DQN formula appears in both target and online network architectures.
private static final BinaryOperator<NDArray> Q_VALUE_MEAN = (value, advantage) -> {
var subMean = advantage.sub(advantage.mean(AXIS_1_ARR, KEEP_DIMS));
return value.add(subMean);
};
Check the max formula below too:
Max implementation:
private static final BinaryOperator<NDArray> Q_VALUE_MAX = (value, advantage) -> {
var subMean = advantage.sub(advantage.max(AXIS_1_ARR, KEEP_DIMS));
return value.add(subMean);
};
Architecture Comparison Between DQN and Dueling DQN
To implement a layer for the Dueling DQN, we extend AbstractBlock, and implement the architecture, where we have one head for the value, and one for the advantage; together you extract the Q-value:
// omitted code ...
public class DuelingLayer extends AbstractBlock {
// omitted code ...
protected final int actions;
protected final BinaryOperator<NDArray> qValueCalc;
protected final SequentialBlock featureBackbone;
protected final SequentialBlock valueHead;
protected final SequentialBlock advantageHead;
private static final BinaryOperator<NDArray> Q_VALUE_MEAN = ...
private static final BinaryOperator<NDArray> Q_VALUE_MAX = ...
// omitted code ...
public DuelingLayer(int actions,
DuelingType duelingType) {
Consumer<SequentialBlock> featureBackboneModifier = featureBackbone -> featureBackbone.add(linear(128))
.add(Activation::relu)
.add(linear(128))
.add(Activation::relu);
Consumer<SequentialBlock> valueModifier = valueHead -> valueHead.add(linear(1));
Consumer<SequentialBlock> advantageModifier = advantageHead -> advantageHead.add(linear(actions));
this(actions, duelingType, featureBackboneModifier, valueModifier, advantageModifier);
}
public DuelingLayer(int actions,
DuelingType duelingType,
Consumer<SequentialBlock> featureBackboneModifier,
Consumer<SequentialBlock> valueHeadModifier,
Consumer<SequentialBlock> advantageHeadModifier) {
super(VERSION);
this.actions = actions;
this.duelingType = duelingType;
this.qValueCalc = switch (duelingType) {
case MEAN -> Q_VALUE_MEAN;
case MAX -> Q_VALUE_MAX;
};
this.featureBackbone = new SequentialBlock();
featureBackboneModifier.accept(featureBackbone);
this.valueHead = new SequentialBlock();
valueHeadModifier.accept(valueHead);
this.advantageHead = new SequentialBlock();
advantageHeadModifier.accept(advantageHead);
addChildBlock(“featureBackbone”, featureBackbone);
addChildBlock(“valueHead”, valueHead);
addChildBlock(“advantageHead”, advantageHead);
}
@Override
protected void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {
featureBackbone.initialize(manager, dataType, inputShapes);
var featureOutShapes = featureBackbone.getOutputShapes(inputShapes);
valueHead.initialize(manager, dataType, featureOutShapes);
advantageHead.initialize(manager, dataType, featureOutShapes);
}
@Override
public NDList forwardInternal(ParameterStore parameterStore, NDList inputs, boolean training, PairList<String, Object> params) {
if (inputs.size() != 1) throw new IllegalArgumentException(“The dueling dqn just accepts one input!”);
var feature = featureBackbone.forward(parameterStore, inputs, training);
var value = valueHead.forward(parameterStore, feature, training).singletonOrThrow();
var advantage = advantageHead.forward(parameterStore, feature, training).singletonOrThrow();
return new NDList(qValueCalc.apply(value, advantage));
}
// omitted code ...
The forwardInternal(…) method implements the Dueling DQN flow.
Now we configure our network to use this new DuelingLayer:
// class DuelingQNetworkMLP
// omitted code ...
public DuelingQNetworkMLP(int observations,
int actions,
Path modelPath,
String prefix,
NDManager parent,
@NonNull DuelingType duelingType) {
this.observations = observations;
this.actions = actions;
this.subManager = subMgr(parent, getClass());
this.model = newModel(getClass(), subManager.getDevice());
this.net = new DuelingLayer(actions, duelingType);
model.setBlock(net);
// omitted code ...
}
// omitted code ...
Then we can run our AgentDuelingDQNMain to see the result (I won’t repeat the code, as it’s very similar to AgentDQNMain).
PER – Prioritized Experience Replay Buffer
The DQN used the Experience Replay Buffer, where at each learned experience, we store the following data:
The problem is that uniform sampling always gives each sample the same probability of being trained on. So if we have difficulty learning a hard transition, it receives the same training frequency as easier and normal transitions, slowing network convergence. Prioritized Experience Replay (PER) assigns a priority to each element to solve this. Since PER introduces priority bias, Importance Sampling corrects it, and we anneal that correction as the network improves over time.
The priority formula is based on the temporal difference error (TD-Error) of the original DQN formula, where for each transition from state s to next state s’, we store this in the buffer.
Check part of the code that contains the priority definition and saving:
// class AgentDQNPER
// omitted code ...
@Override
protected float trainOnline(int batchSize, IReplayBuffer ireplayBuffer, Loss lossFunc, NDManager sub) {
// omitted code ...
perl2LossFunc.normISWeights(samples.weights());
// (perl2LossFunc) L = sum norm w^{is} * error^2
@Cleanup var losses = rawBackwardLoss(sub, perl2LossFunc, targetQValue, () -> {
var states = samples.states();
var actions = samples.actions();
// y_hat = q_online(s, a)
return onlineNet.forward(states, qValue -> qValue.gather(actions, AXIS_1));
});
var lossItem = scopedToFloat(NDArray::mean, losses);
@Cleanup var priorities = scoped(it -> it.abs().add(MIN_PRIORITY), losses);
replayBuffer.updatePriorities(samples.bufferIndexes(), priorities);
// omitted code ...
}
// omitted code ...
We add the transition in the buffer, so it can be used:
In the code below the defaultPriority starts with one, but it’s adjusted when we update the priorities (the update will be shown later):
// class PrioritizedReplayBuffer
// omitted code ...
@Override
public void store(Experience exp) {
exp.state().attach(subManager);
exp.nextState().attach(subManager);
if (size < capacity) ++size;
var oldExp = experiences[pos];
experiences[pos] = exp;
float defaultPriority = (float) Math.pow(maxPriority, alpha);
sumSegmentTree.update(pos, defaultPriority);
minSegmentTree.update(pos, defaultPriority);
if (oldExp != null) oldExp.close();
pos = (pos + 1) % experiences.length;
}
// omitted code ...
Inside PER, we use a Segment Tree because this type of data structure gives us O(log n) complexity for updating the priority in the tree and O(1) complexity when getting the total sum (if it is a Sum Segment Tree). To efficiently scale PER operations when you have thousands (or more) of transitions in the buffer, we use both a Sum Segment Tree and a Min Segment Tree. The Sum Segment Tree is used to sample transitions, while the Min Segment Tree is used to get the minimum priority, which is used for Importance Sampling weight normalization.
Segment Tree
Check the example of Sum Segment Tree (almost the same as Min Segment Tree), using the tree data structure (in the implementation we have the flat array too):
When you need to update, you check if the value being searched is greater or lower than the sum. If it’s lower, we just go to the left node; if it’s bigger, we go to the right by subtracting the left side. Why do we subtract the left side? If we didn’t do that, we would instruct the algorithm that the “map” remains completely accessible, which is a bug. To correct this, we must subtract to instruct that the other “part” of the “map” no longer remains accessible, reducing the search area. After that, we update the node and go up in O(log n), updating the parent nodes.
Range Sum Operation
In the rangeSum algorithm we descend just to the level of the index that we found; if it match, it is not necessary to descend anymore. Check the image below:
Check the implementation below:
//class SumSegmentTree
// omitted code ..
private float rangeSum(Node node, int start, int end) {
if (start > node.rightMost || end < node.leftMost) return 0;
if (start <= node.leftMost && end >= node.rightMost) return node.sum;
return rangeSum(node.left, start, end) + rangeSum(node.right, start, end);
}
// omitted code ..
Sample Index by Value: Navigation Logic
In the sampleIndexByValue(...) method, we search for any random value whose index falls between the start and end part of each segment. The PER algorithm generates the segments by dividing the total value by the size of the mini-batch. The navigation follows the same logic as the update operation: we descend to the left node when going left, but to go right, we must subtract the left value.
Operation sampleIndexByValue(sample=0.55) – navigates the tree using cumulative sums. (Source: Author’s own work, assisted by AI / generated via Claude AI and verified by the author)The operations sampleIndexByValue and rangeSum are not necessary in Min Segment Tree.
Another perspective on why we must subtract the left values when going to the right in the sampleIndexByValue operation: for example, if we want to access the block where the probability is closest to 0.55, we go right, but if we don’t subtract the left sum, the algorithm would land on [3, 3] instead of [2, 2], which is the correct target.
Implementation
Below you can see the implementation of this method:
// class SumSegmentTree
// omitted code ..
@Override
public int sampleIndexByValueInRange(float lower, float upper) {
float totalSum = sum();
if (lower < 0) throw new IllegalArgumentException("lower cannot be negative: " + lower);
if (upper > totalSum) throw new IllegalArgumentException("upper cannot exceed total sum. upper: " + upper + ", total: " + totalSum);
if (lower >= upper) throw new IllegalArgumentException("lower must be < upper. lower: " + lower + ", upper: " + upper);
float sample = random.nextFloat(lower, upper);
return sampleIndexByValue(root, sample);
}
private int sampleIndexByValue(Node node, float sample) {
if (node.leftMost == node.rightMost) return node.leftMost;
if (sample <= node.left.sum) return sampleIndexByValue(node.left, sample);
return sampleIndexByValue(node.right, sample - node.left.sum);
}
// omitted code ..
Prioritized Experience Replay Buffer
Now that the Segment Tree is more clear, we can understand in depth how each formula is connected. The P(i) is the sampling of priority transitions, which is saved in the Sum and Min Segment Trees, where the bigger the P(i) the larger the segment it occupies, meaning that it will be sampled more frequently. The alpha shows that, if it’s 0, all priorities would have the same size, and the segments would be of the same size too. Check the formulas below:
In the prioritizedIndexSamples we divide the areas between segments; remember that we can sample the same transition many times; because of this, Importance Sampling is applied (more about Importance Sampling later):
// class PrioritizedReplayBuffer
// omitted code ...
protected int[] prioritizedIndexSamples(int batchSize) {
var batchIndexes = new int[batchSize];
float segment = sumSegmentTree.sum() / batchSize;
for (int i = 0; i < batchSize; ++i) {
float lower = i * segment;
float upper = (i + 1) * segment;
int sampledIdx = sumSegmentTree.sampleIndexByValueInRange(lower, upper);
batchIndexes[i] = sampledIdx;
}
return batchIndexes;
}
// omitted code ...
public VecExperience sample(int batchSize, float beta) {
if (!enough(batchSize)) return null;
var sub = subMgr(subManager, “prioritized-sample”);
var bufferIndexes = prioritizedIndexSamples(batchSize);
var batch = buildPrioritizedSamples(bufferIndexes);
var states = newAttachedList(sub, batch, exp -> exp.state().duplicate());
var actions = djlMapToLong(sub, batch, exp -> exp.actionAs(Long.class));
var rewards = djlMapToFloat32(sub, batch, Experience::reward);
var nextStates = newAttachedList(sub, batch, exp -> exp.nextState().duplicate());
var dones = djlMapToFloat32(sub, batch, exp -> exp.done() ? 1 : 0);
var weights = calculateWeights(sub, bufferIndexes, beta);
return new VecExperience(
sub,
states,
actions,
rewards,
nextStates,
dones,
weights,
bufferIndexes
);
}
// omitted code ...
And now we have the definition of the Importance Sampling weights formula, used to correct the bias introduced by the priority sampling; without this, the network would only focus on learning the hard samples and would overfit. In the next graphs, we’ll have a textual and visual explanation of how they work:
Code of the Importance Sampling:
// class PrioritizedReplayBuffer
// omitted code ...
private NDArray calculateWeights(NDManager sub, int[] bufferIndexes, float beta) {
float sum = sumSegmentTree.sum();
float pMin = Math.max(minSegmentTree.min() / sum, MIN_PRIORITY);
int count = minSegmentTree.size();
float maxISWeight = (float) Math.pow(count * pMin, -beta);
return sub.create(stream(bufferIndexes)
.mapToDouble(idx -> {
float pSample = sumSegmentTree.get(idx) / sum;
return Math.pow(count * pSample, -beta) / maxISWeight;
})
.toArray())
.toType(DataType.FLOAT32, false)
.expandDims(1)
.toDevice(device, false);
}
// omitted code ...
The graphic below shows that Pmin has a weight of 1, meaning that its value unscaled in the backpropagation; the other weights, which have greater priority, will have a smaller weight, meaning that their value will be downscaled proportionally to the magnitude of their priority.
You could calculate these weights by position too, or put them manually: for the first position, weight 1, for the second, weight 0.8, and so on. Of course it would not be efficient, but we have a less efficient but good enough algorithm called rank that calculates the priority based on the position in the rank, like this: 1/rank(i).
Another graph showing the weight value reduces and the priority increases:
PER Loss Function
The first formula is the TD-Error, and the styled L is the loss used in Backpropagation where each error is scaled by the normalized weight; note that we are using MSE (Mean Square Error), but you can use the Huber Loss10, MAE (Mean Absolute Error)11 and others in their place:
Code of the loss function. Note that by default we don’t apply mean in the output of evaluate(…) function, because we need to get the loss per element:
// omitted code ...
public class PERL2Loss extends Loss {
private static final float HALF_WEIGHT = 0.5f;
private final Reduction reduction;
private NDArray normISWeights;
// omitted code ...w PERL2Loss(Reduction.NONE);
public PERL2Loss normISWeights(NDArray normISWeights) {
if (normISWeights.getShape().dimension() != 2) {
normISWeights = normISWeights.reshape(N_BATCH, 1);
}
this.normISWeights = normISWeights;
return this;
}
@Override
public NDArray evaluate(NDList label, NDList prediction) {
if (null == normISWeights) {
throw new IllegalStateException("You must set the normISWeights before training!");
}
var pred = prediction.singletonOrThrow();
var labelReshaped = label.singletonOrThrow().reshape(pred.getShape());
var timeDifferenceError = labelReshaped.sub(pred).square();
var weightedError = timeDifferenceError.mul(normISWeights);
var loss = weightedError.mul(HALF_WEIGHT);
this.normISWeights = null;
if (reduction.equals(Reduction.MEAN)) {
return loss.mean();
}
// (batch) - We need to extract the loss per element to update sum/min segment-tree
return loss.squeeze();
}
}
Training Method with PER and Loss Integration
Training method of AgentDQNPER:
// class AgentDQNPER
// omitted code ...
@Override
protected float trainOnline(int batchSize, IReplayBuffer ireplayBuffer, Loss lossFunc, NDManager sub) {
if (!ireplayBuffer.enough(batchSize)) return Float.NaN;
if (!(ireplayBuffer instanceof PrioritizedReplayBuffer replayBuffer)) {
throw new IllegalArgumentException(“You must pass PrioritizedReplayBuffer!”);
}
if (!(lossFunc instanceof PERL2Loss perl2LossFunc)) {
throw new IllegalArgumentException(“You must pass PERL2Loss!”);
}
@Cleanup var samples = replayBuffer.sample(batchSize, beta);
@Cleanup var targetQValue = targetNet.forward(samples.nextStates(), nextQValue -> {
// max Q(s’, a’)
var maxNextQValue = nextQValue.max(AXIS_1_ARR, KEEP_DIMS);
// gamma * max Q(s’, a’)
var discountNextQValue = maxNextQValue.mul(gamma);
// (1 - done)
var mask = samples.dones().neg().add(1);
// r + gamma * max Q(s’, a’) * (1 - done)
return samples.rewards()
.add(discountNextQValue.mul(mask))
.stopGradient();
});
perl2LossFunc.normISWeights(samples.weights());
// (perl2LossFunc) L = sum norm w^{is} * error^2
@Cleanup var losses = rawBackwardLoss(sub, perl2LossFunc, targetQValue, () -> {
var states = samples.states();
var actions = samples.actions();
// y_hat = q_online(s, a)
return onlineNet.forward(states, qValue -> qValue.gather(actions, AXIS_1));
});
var lossItem = scopedToFloat(NDArray::mean, losses);
@Cleanup var priorities = scoped(it -> it.abs().add(MIN_PRIORITY), losses);
replayBuffer.updatePriorities(samples.bufferIndexes(), priorities);
trainStep(onlineNet.getBlock(), optimizer);
return lossItem;
}
// omitted code ...
Priority Update Mechanism
The code below shows how we update the priorities after each training.
// class PrioritizedReplayBuffer
// omitted code ...
public void updatePriorities(int[] bufferIndexes, NDArray priorities) {
if (bufferIndexes.length != priorities.size()) throw new IllegalArgumentException(“Invalid length!”);
float tempMaxPriority = this.maxPriority;
var rawPriorities = toFloatArray(priorities);
for (int i = 0; i < bufferIndexes.length; ++i) {
int experienceIdx = bufferIndexes[i];
if (experienceIdx < 0 || experienceIdx >= experiences.length) {
throw new ArrayIndexOutOfBoundsException(“Invalid experienceIdx: “ + experienceIdx);
}
float priority = (float) Math.pow(rawPriorities[i], this.alpha);
if (priority <= 0) {
throw new IllegalArgumentException(“priority must be greater than 0!”);
}
this.sumSegmentTree.update(experienceIdx, priority);
this.minSegmentTree.update(experienceIdx, priority);
tempMaxPriority = Math.max(tempMaxPriority, priority);
}
this.maxPriority = tempMaxPriority;
}
// omitted code ...
Alpha and Beta Annealing in PER
The animated graph below shows the alpha (prioritization) and beta (IS correction) are being annealed over time, where they are the “opposite” of each other in this context:
Alpha annealing is not always used; it depends on the implementation. The original article, “Prioritized Experience Replay”12, maintains the alpha at 0.6, while only beta is annealed over time.
Importance-Sampling Weight Correction in Practice
Below there’s another table showing, with concrete values, how the PER works with the Importance-Sampling Weight correction. You can see that the bigger the error magnitude, the smaller the weight we attribute to this sampling. This avoids biasing the neural network toward learning only hard samples, since we train them much more than the other samples. You can see that sample A can train almost 8 times more than sample D. Note that because we rarely see sample D, its weight is 1; we don’t increase it above 1 because that would introduce another type of bias:
Biasing the network toward learning hard samples can make it overfit on them without generalizing well. That is why we anneal beta over time; as the network trains, the error reduces and the distribution tends toward uniform. It will not become fully uniform if we fix alpha with a value greater than 0.
N-Step Learning
N-step learning is a middle ground between the TD-Error and Monte Carlo (MC) methods13, where in MC, you go through the episode until the end to get the best estimation of the reward. You can see below in the formula that we sum until we reach n plus a bootstrap term:
Next, you can make comparisons between TD-Error and N-step with values of 2, 3, and 4:
The difference between TD-Error, N-step, and Monte Carlo is fundamentally about a trade-off between bias and variance.
The 3-step approach is used in Rainbow DQN14
Algorithm Overview
In the algorithm below we accumulate the transitions n times, and when we reach the last transition (n = 3 for example), we do the following:
- State/action: We get from the first experience object (
n-step = 1) the state asoldStateandactionasoldAction. - Rewards: All rewards (from
n-step 1,2, and3) will be accumulated using the above formula intonStepReward. - Next state: The next state of the last transition (
n-step = 3) will benewNextState.
These new values – oldState, oldAction, nStepReward, newNextState, and done; will be stored in the parent class, which can be ExperienceReplayBuffer or PrioritizedReplayBuffer (if it’s NStepPrioritizedReplayBuffer).
If done = true when n-step less than or equal to 3, for example, we repeat the same process, with the difference that we have to remove all transitions that reach the terminal state (done = true); if we don’t remove them, we may have an invalid transition, like a done pointing to an undone transition.
Diagram:
Java Implementation
Check the code below:
// omitted code ...
public class NStepExperienceReplayBuffer extends ExperienceReplayBuffer {
private final int nStep;
private final float gamma;
private final Deque<Experience> nStepDeque;
// omitted code ...
public NStepExperienceReplayBuffer(int capacity,
int nStep,
float gamma,
NDManager parent,
ExperienceSampler sampler) {
super(capacity, parent, sampler);
if (nStep < 1) throw new IllegalArgumentException(“nStep must at least 1: “ + nStep);
if (gamma <= 0 || gamma > 1) throw new IllegalArgumentException(“gamma must be between range (0, 1]: “ + gamma);
this.nStep = nStep;
this.gamma = gamma;
this.nStepDeque = new ArrayDeque<>(nStep);
}
@Override
public void store(Experience exp) {
nStepDeque.addLast(exp);
if (!exp.done()) {
if (nStepDeque.size() == nStep) {
emitExperience();
nStepDeque.pollFirst().close();
}
return;
}
while (!nStepDeque.isEmpty()) {
emitExperience();
nStepDeque.pollFirst().close();
}
}
private void emitExperience() {
double nStepReward = 0f;
var oldest = nStepDeque.getFirst();
int k = 0;
Experience newest = null;
for (var curr : nStepDeque) {
nStepReward += Math.pow(gamma, k) * curr.reward();
newest = curr;
if (curr.done()) break;
++k;
}
requireNonNull(newest, “Cannot be null!”);
var oldState = oldest.state().duplicate();
var oldAction = oldest.action().duplicate();
var newNextState = newest.nextState().duplicate();
var done = newest.done();
var nStepExperience = new Experience(oldState, oldAction,
nStepReward, newNextState, done);
super.store(nStepExperience);
}
public int nStep() {
return nStep;
}
@Override
public void close() {
super.close();
nStepDeque.forEach(Experience::close);
}
}
Then we finish with the DQN algorithm applied with n-step learning:
// class AgentNStepDQN
// omitted code ...
@Override
protected float trainOnline(int batchSize, IReplayBuffer ireplayBuffer, Loss lossFunc, NDManager sub) {
if (!ireplayBuffer.enough(batchSize)) return Float.NaN;
if (!(ireplayBuffer instanceof NStepExperienceReplayBuffer replayBuffer)) {
throw new IllegalArgumentException(“You must pass NStepExperienceReplayBuffer!”);
}
@Cleanup var samples = replayBuffer.sample(batchSize);
@Cleanup var targetQValue = targetNet.forward(samples.nextStates(), nextQValue -> {
var rewards = samples.rewards();
var dones = samples.dones();
// max Q(s’, a’)
var maxNextQValue = nextQValue.max(AXIS_1_ARR, KEEP_DIMS);
float gammaNBootstrap = (float) Math.pow(gamma, replayBuffer.nStep());
// gamma^n * max Q(s’, a’)
var discountNextQValue = maxNextQValue.mul(gammaNBootstrap);
// (1 - done)
var mask = dones.neg().add(1);
// n-step reward + gamma^n * max Q(s', a') * (1 - done)
return rewards
.add(discountNextQValue.mul(mask))
.stopGradient();
});
float lossItem = backwardLoss(sub, lossFunc, targetQValue, () -> {
var states = samples.states();
var actions = samples.actions();
// y_hat = q_online(s, a)
return onlineNet.forward(states, qValue -> qValue.gather(actions, AXIS_1));
});
trainStep(onlineNet.getBlock(), optimizer);
return lossItem;
}
// omitted code ...
Note the variable discountNextQValue – it’s the bootstrap term.
Another thing you may notice is that we can use many blocks together, like N-step learning with Double DQN, Dueling DQN, and PER, for example.
Noisy Networks
When we train using greedy selection, a technique where we handle exploration outside of the network, we need to manually select the correct epsilon, with a value that decreases over enough time so the network can learn.
This stands in contrast to Noisy Networks, where we build a network that learns how to filter noise and learns about the environment, making the traditional epsilon-greedy method unnecessary.
Let’s start overriding the selectAction(…) method:
// omitted code ...
public class AgentNoisyNetDQN extends AbstractAgent {
// omitted code ...
@Override
public ActionSpaceType.ActionResult selectAction(NDArray state) {
onlineNoisyNet.resetNoise();
@Cleanup var oneBatchState = state.expandDims(0);
long action = onlineNoisyNet.forwardLong(oneBatchState,
qValue -> qValue.stopGradient().argMax(AXIS_1));
return actionSpaceType.get(action);
}
}
Noise Distributions and the Reparameterization Trick
This noise is created using the Normal Distribution15. Check the graphic below:
Each one of these curves is a possible distribution for the weights and biases, updated via backpropagation through the network. We shift the distribution with the μ (mean), and control exploration with σ (standard deviation), where if it’s wider (σ = 2 or σ = 3), then the exploration tends to be more broad, and if it’s more concentrated in the middle (σ = 0.4 or σ = 0.36), it’s more greedy.
To train the network using this noise as perturbation, and explore more in the beginning and exploit more in the end automatically, we need to do the reparameterization trick, which is similar to the one applied to VAE (Variational Autoencoder)16. You can’t backpropagate through a sample, and if you attach this noise generator to the network, the network will treat the noise as a weight and will no longer generate the values for the distribution. The following image illustrates this process:
Note that it’s applied in other machine learning techniques too17.
Below, we display the expansion formula used during training:
At inference time, where it’s not necessary to generate samples, the epsilon is zeroed. Follow the arrows to see the simplification of the formula (only at inference/test time):
Custom Implementation in DJL
The code below uses a different structure for a neural network in DJL, as we are working with a custom workflow and execution, like sampling a Gaussian distribution. We configure the weights and biases as parameters and use them in the flow forwardInternal(…):
// omitted code ...
public class NoisyLayer extends AbstractBlock implements AutoCloseable {
// omitted code ...
private final int outFeatures;
private final Parameter weightMu;
private final Parameter weightSigma;
private final Parameter biasMu;
private final Parameter biasSigma;
private FactorizedNoise.Noise noise;
public NoisyLayer(int outFeatures) {
super(VERSION);
this.outFeatures = outFeatures;
this.weightMu = addParameter(Parameter.builder()
.setName(“weightMu”)
.setType(Parameter.Type.WEIGHT)
.build());
this.weightSigma = addParameter(Parameter.builder()
.setName(“weightSigma”)
.setType(Parameter.Type.WEIGHT)
.build());
this.biasMu = addParameter(Parameter.builder()
.setName(“biasMu”)
.setType(Parameter.Type.BIAS)
.build());
this.biasSigma = addParameter(Parameter.builder()
.setName(“biasSigma”)
.setType(Parameter.Type.BIAS)
.build());
}
@Override
protected NDList forwardInternal(ParameterStore paramStore,
NDList inputs,
boolean training,
PairList<String, Object> params) {
@Cleanup var sub = subMgr(inputs.getManager(), “noisy-layer”);
var device = inputs.getManager().getDevice();
var input = inputs.singletonOrThrow();
var wMu = paramStore.getValue(weightMu, device, training);
var wSigma = paramStore.getValue(weightSigma, device, training);
var bMu = paramStore.getValue(biasMu, device, training);
var bSigma = paramStore.getValue(biasSigma, device, training);
sub.tempAttachAll(wMu, wSigma, bMu, bSigma);
if (training) {
int inFeatures = (int) input.getShape().getLastDimension();
ensureNoiseIsSampled(input.getManager(), inFeatures, outFeatures);
// w = μʷ + σʷ ⊙ εʷ
var w = wMu.add(wSigma.mul(noise.epsWeight()));
// b = μᵇ + σᵇ ⊙ εᵇ
var b = bMu.add(bSigma.mul(noise.epsBias()));
// y = w · x + b, in ML is x · w + b
return Linear.linear(input, w, b);
}
return Linear.linear(input, wMu, bMu);
}
// omitted code ...
}
We use addParameter(…) to register trainable parameters with DJL. This enables weight initialization, parameter tracking for our custom DJLOptimizer (which handles trainStep(…) and gradient clipping), checkpoint saving/loading (feature in progress), and (in principle) multi-device synchronization if we ever move from the current single-device setup. We avoid the built-in Trainer class because our custom optimizer gives us more control over the training loop, and the Trainer implementation is designed for a single Model, and we use two models, online and target, for DQN implementations. We don’t use the gradient validation feature either.
Factorized Noise: Efficient Sampling via Outer Product
Then, how do we generate this epsilon efficiently? With the formula below, we can sample a distribution similar to the normal distribution but with lower computational cost:
Check the code:
// omitted code ...
public class FactorizedNoise {
private FactorizedNoise() {
throw new IllegalStateException(“No FactorizedNoise!”);
}
// omitted documentation ...
static NDArray sampleNoise(NDManager manager, int size) {
var gaussian = manager.randomNormal(0f, 1f, new Shape(size), DataType.FLOAT32);
return gaussian.sign().mul(gaussian.abs().sqrt());
}
public static Noise sampleNoiseOuter(NDManager manager, int inSize, int outSize) {
@Cleanup var sub = subMgr(manager, “samp-noise-outer”);
var fepsOut = sampleNoise(sub, outSize);
var fepsIn = sampleNoise(sub, inSize);
// (out x 1) x (1 x in) = (out x in)
var epsW = fepsOut.reshape(outSize, 1)
.matMul(fepsIn.reshape(1, inSize));
return new Noise(
// epsWeight = f(epsOut) * f(epsIn)^T
sub.ret(epsW),
// epsBias = f(epsOut)
sub.ret(fepsOut)
);
}
public record Noise(NDArray epsWeight,
NDArray epsBias) implements AutoCloseable {
@Override
public void close() {
epsWeight.close();
epsBias.close();
}
public boolean isRelease() {
return epsWeight.isReleased() && epsBias.isReleased();
}
}
}
Verifying Noise Samples Are Not Trainable
And as we affirmed above, the epsilon noise samples are temporary and we never backpropagate them. The test below shows this:
// class FactorizedNoiseTest
// omitted code ...
@Test
void shouldSampleNoiseOuterBeUntrainable() {
var expectedRequiresGrad = false;
@Cleanup var noise = FactorizedNoise.sampleNoiseOuter(manager, 7, 14);
assertEquals(expectedRequiresGrad, noise.epsWeight().hasGradient());
assertEquals(expectedRequiresGrad, noise.epsBias().hasGradient());
}
// omitted code ...
Step-by-Step Sample Generation Process
Next, we can see step by step how the samples are generated. Check the table below, where we use vectors with 2 values each as an explanation:
We apply the function to both vectors:
Now we generate the weight matrix by applying the outer product between the two vectors:
Then we do the same with the bias:
Now the process is done.
Efficiency Comparison: Factorized vs. Direct Sampling
Below we compare the number of operations when sampling using the Normal Distribution directly versus generating the factorized matrix via the outer product. As you can see, the outer product is more computationally efficient, and the authors mention that even though the generated matrix is dependent, there is not too much loss in training quality.
Computational Efficiency: Factorized vs. Independent Sampling
The graphic below shows when this efficiency gain between factorized and independent sampling (using a Normal Distribution generator) pays off in computational efficiency.
Factorized noise samples p + q scalars and applies f(x) = sign(x)√|x| before the outer product. The fixed overhead dominates at small sizes but amortizes rapidly; independent grows as O(p·q) while factorized stays near O(p+q)18.
Below you can see graphically the difference between Factorized noise (using the outer product between two vectors) and independent sampling (Normal Distribution generator).
Note that in the factorized heatmaps (first column) the pattern is more structured and predictable compared to independent sampling (fourth column), and the factorized noise histogram (second column) is bimodal, an effect of the f(x)=sign(x)√∣x∣ transformation, while the independent sampling histogram (fifth column) is unimodal.
Initialization Strategy for Noisy Networks
How do we initialize the noisy networks? The papers suggest the following patterns:
We do this to normalize the weights’ initialization, which lets variance stay constant independent of the number of inputs and avoids gradient exploding (or vanishing). Check the graphic below:
Noisy Initialization Implementation
The NoisyLayerInit, as the name suggests, is responsible for initializing the NoisyLayer, and later, the NoisyLayer class can be used many times by the NoisyQNetworkMLP and its derivatives. See below:
// omitted code ...
public class NoisyLayerInit implements Initializer {
private static final float BASE_INIT = 1f;
private static final float SIGMA_INIT = 0.5f;
public enum InitType {
MU, SIGMA
}
public static NoisyLayerInit ofMu(int size) {
return new NoisyLayerInit(size, InitType.MU);
}
public static NoisyLayerInit ofSigma(int size) {
return new NoisyLayerInit(size, InitType.SIGMA);
}
@Getter
private final int size;
@Getter
private final InitType type;
private NoisyLayerInit(int size, InitType type) {
this.size = size;
this.type = type;
}
@Override
public NDArray initialize(NDManager manager, Shape shape, DataType dataType) {
float p = (float) Math.sqrt(size);
return switch (type) {
case MU -> {
float muRange = BASE_INIT / p;
yield manager.randomUniform(-muRange, muRange, shape, dataType);
}
case SIGMA -> manager.full(shape, SIGMA_INIT/p, dataType);
};
}
}
Then see how it’s used in NoisyLayer. Note the NoisyLayerInit ofMu(…) and ofSigma(…) methods:
// omitted code ...
public class NoisyLayer extends AbstractBlock implements AutoCloseable {
// omitted code ...
@Override
public void prepare(Shape[] inputShapes) {
int inFeatures = (int) inputShapes[0].getLastDimension();
var weightShape = new Shape(outFeatures, inFeatures);
var biasShape = new Shape(outFeatures);
weightMu.setShape(weightShape);
weightSigma.setShape(weightShape);
biasMu.setShape(biasShape);
biasSigma.setShape(biasShape);
// muW, muB ~ U(-1 / sqrt(p_in), 1 / sqrt(p_in))
weightMu.setInitializer(NoisyLayerInit.ofMu(inFeatures));
biasMu.setInitializer(NoisyLayerInit.ofMu(inFeatures));
// sigmaW = 0.5 / sqrt(p_in)
weightSigma.setInitializer(NoisyLayerInit.ofSigma(inFeatures));
// sigmaB = 0.5 / sqrt(p_out) -> it’s asymmetric
biasSigma.setInitializer(NoisyLayerInit.ofSigma(outFeatures));
}
// omitted code ...
}
Building the Noisy Network and Training Loop
Note the noisyLayers attribute addition and noisyLayer(…) invocation to create the noisy networks on the net variable:
// omitted code ...
public class NoisyQNetworkMLP implements INoisyNetwork {
private boolean training;
private final NDManager subManager;
private final int observations;
private final int actions;
private final Model model;
private final SequentialBlock net;
private final List<NoisyLayer> noisyLayers;
private final ParameterStore parameterStore;
// omitted code ...
@SneakyThrows
public NoisyQNetworkMLP(int observations,
int actions,
Path modelPath,
String prefix,
NDManager parent) {
this.observations = observations;
this.actions = actions;
this.subManager = subMgr(parent, getClass());
this.noisyLayers = new ArrayList<>();
this.model = newModel(getClass(), subManager.getDevice());
this.net = new SequentialBlock();
net.add(linear(128))
.add(Activation::relu)
.add(addAndGet(noisyLayers, noisyLayer(128)))
.add(Activation::relu)
.add(addAndGet(noisyLayers, noisyLayer(actions)));
model.setBlock(net);
// omitted code ...
}
@Override
public void resetNoise() {
noisyLayers.forEach(NoisyLayer::resetNoise);
}
// omitted code ...
}
We call resetNoise() before each training step to refresh the noise samples in AgentNoisyNetDQN. Without this, the network would reuse the same noise samples and eventually overfit.
// class AgentNoisyNetDQN
// omitted code ...
@Override
protected float trainOnline(int batchSize, IReplayBuffer replayBuffer, Loss lossFunc, NDManager sub) {
if (!replayBuffer.enough(batchSize)) return Float.NaN;
onlineNoisyNet.resetNoise();
@Cleanup var samples = replayBuffer.sample(batchSize);
@Cleanup var targetQValue = targetNoisyNet.forward(samples.nextStates(), nextQValue -> {
// max Q(s’, a’)
var maxNextQValue = nextQValue.max(AXIS_1_ARR, true);
// gamma * max Q(s’, a’)
var discountNextQValue = maxNextQValue.mul(gamma);
// (1 - done)
var mask = samples.dones().neg().add(1);
// r + gamma * max Q(s’, a’) * (1 - done)
return samples.rewards()
.add(discountNextQValue.mul(mask))
.stopGradient();
});
float lossItem = backwardLoss(sub, lossFunc, targetQValue, () -> {
var states = samples.states();
var actions = samples.actions();
// y_hat = q_online(s, a)
return onlineNoisyNet.forward(states, qValue -> qValue.gather(actions, AXIS_1));
});
trainStep(onlineNoisyNet.getBlock(), optimizer);
return lossItem;
}
// omitted code ...
Take a look at AgentNoisyDuelingNetDDQN as well; this version brings together several of the concepts discussed in this article.
Shannon Entropy and Kullback-Leibler Intuition
Before introducing Categorical DQN19, I want to explain Shannon Entropy conceptually and show how it can be applied at the end of this section.
Let’s define the simple system below, where nodes have transition probabilities to other nodes.
Calculating the Shannon entropy of this system yields 0.4406 bits. You can see the step-by-step calculation at this link.
Next, we add new transitions to a separate partition of the system to compare the two. See the image below:
Calculating the Shannon entropy of System 2 we get 0.6752 bits; you can see the step-by-step calculation at this link.
Comparing the entropy between them, we see that H(System 2) > H(System 1), where 0.6752 > 0.4406. This shows that the uncertainty of System 2 is greater than that of System 1, meaning System 2 is more complex.
Now, we want to check how much information we would lose by replacing System 1 with System 2. Calculating the Kullback-Leibler Divergence, we get 0.1969 bits (check this link).
The 0.1969 bits in KL divergence doesn’t mean the new system is worse, but that information is lost. In a real production system, you would need to retrain the model to match this new distribution. Or, as I’ll explain later in Categorical DQN, the online Q-prediction will be retrained online with the goal of minimizing KL, which reduces this loss of information.
Visualizing the KL Divergence Between Systems
Below, you can see in the figure where P is greater than Q and where Q is greater than P. The E node contributes zero because it doesn’t exist in the P distribution.
In case you want to deepen your understanding, check the link in footnote20.
Now, let’s proceed to the C51 algorithm.
Distributional DQN (C51)
C51 works as a Distributional Reinforcement Learning method where we structure returns as a probability distribution, instead of treating state-action values as scalars. For example, jumping over the pit would have a distribution shifted more to the right than going straight and falling into the pit, which would have a distribution almost entirely to the left (where the character dies in the game). Check an example of the distribution:
See below an example of states and actions:
With a distribution, the network can optimize more efficiently during training because we can shift it (left or right) to match the target distribution.
See the graph below showing the steps for the Bellman projection. Using this technique, we can incorporate the reward and discount factor (gamma) into the distribution and make the network learn:
Original Distributions
The first step is to calculate the difference in value (delta z) for each position in the support vector. Then, we create the support vector (z) and fill it with default values:
Bellman Projection
Now we shift the values for each position in the support vector and clip them (this entire process is the Bellman projection), using b as a continuous index that maps to the lower atom (l) or upper atom (u):
Now we redistribute the probability mass to its neighboring atoms:
Bellman Projection Java Implementation
Entire code of Bellman projection:
// omitted code ...
@Accessors(fluent = true)
public class CategoricalBellmanProjection {
public static final int N_ATOMS = 51;
public static final int SELECTED_ACTION_PER_BATCH = 1;
public static final float V_MIN = -10.0f;
public static final float V_MAX = +10.0f;
@Getter
private final int atoms;
@Getter
private final float vMin;
@Getter
private final float vMax;
@Getter
private final float zDelta;
private final float[] supportVectorZ;
private final Shape expectedShape;
public CategoricalBellmanProjection() {
this(N_ATOMS, V_MIN, V_MAX);
}
public CategoricalBellmanProjection(int atoms,
float vMin,
float vMax) {
this.atoms = atoms;
this.vMin = vMin;
this.vMax = vMax;
this.zDelta = (vMax - vMin)/(atoms - 1);
this.supportVectorZ = generateSupportVectorZ(atoms, vMin, zDelta);
this.expectedShape = new Shape(N_BATCH, 1, atoms);
}
public float[] support() {
return supportVectorZ.clone();
}
public NDArray support(NDManager mgr) {
return mgr.create(supportVectorZ);
}
private float[] generateSupportVectorZ(int atoms, float vMin, float zDelta) {
var zSupport = new float[atoms];
for (int i = 0; i < atoms; ++i) {
zSupport[i] = vMin + i * zDelta;
}
return zSupport;
}
public NDArray project(final NDArray probNextDist,
final NDArray rewards,
final NDArray dones,
final float gamma) {
var inputShape = probNextDist.getShape();
if (!equalsDims(inputShape, expectedShape, 1)) {
throw new IllegalArgumentException(“Invalid shape “ + inputShape + “, must be (batch, 1, %d) “.formatted(atoms));
}
if (inputShape.get(0) != rewards.getShape().get(0) || inputShape.get(0) != dones.getShape().get(0)) {
throw new IllegalArgumentException(“Invalid dimensions between probNextDist (len=%d), rewards (len=%d) or dones (len=%d)”
.formatted(inputShape.get(0), rewards.getShape().get(0), dones.getShape().get(0)));
}
@Cleanup var sub = subMgr(probNextDist, “bellman-proj”);
sub.tempAttachAll(probNextDist, rewards, dones);
int batchSize = (int) rewards.getShape().get(0);
float[] rewardsArr = rewards.reshape(FLATTEN).toFloatArray();
float[] probNextDistArr = probNextDist.reshape(FLATTEN).toFloatArray();
float[] donesArr = dones.toFloatArray();
float[] massDist = new float[batchSize * atoms];
for (int batch = 0; batch < batchSize; ++batch) {
int offset = batch * atoms;
float reward = rewardsArr[batch];
float notDone = 1 - donesArr[batch];
for (int atomIdx = 0; atomIdx < atoms; ++atomIdx) {
// Tz = r + γ * z_j * (1 - done)
float targetProj = Math.clamp(reward + gamma * supportVectorZ[atomIdx] * notDone, vMin, vMax);
// b = (Tz - vMin) / ∆z
float baseIdx = (targetProj - vMin) / zDelta;
int lowerIdx = (int) Math.floor(baseIdx);
int upperIdx = (int) Math.ceil(baseIdx);
float prob = probNextDistArr[offset + atomIdx];
// m_l += p(s’, a, θ-) * (u - b)
// m_u += p(s’, a, θ-) * (b - l)
if (lowerIdx == upperIdx) {
massDist[offset + lowerIdx] += prob;
} else {
massDist[offset + lowerIdx] += prob * (upperIdx - baseIdx);
massDist[offset + upperIdx] += prob * (baseIdx - lowerIdx);
}
}
}
var massDistTarget = sub.create(massDist, new Shape(batchSize, SELECTED_ACTION_PER_BATCH, atoms));
return sub.ret(massDistTarget);
}
}
It’s implemented using pure Java because DJL didn’t have the scatterAddi method in its API. If/when they add it, we can modify the code, and this may improve performance as we can execute on the GPU.
Impact of Atom Count on Distribution Resolution
Now, with the target network’s distribution projected onto m, we update the online network:
The image below shows how varying the number of atoms impacts resolution; a higher atom count yields a smoother probability distribution, allowing the network to capture richer environment dynamics:
When you add the reward and apply the discount factor to the support, the distribution changes – as shown in the plots below, separated by the number of atoms:
Support transformation behavior for positive, zero, and negative reward values:
Below, you can see the architecture defined with 3, 5, 10, and 25 atoms:
Note that as the number of atoms increases, the number of parameters in your network will increase, which may increase accuracy, but at a computational cost.
Network Implementation in DJL
Note the main modifications below:
- The output layer in the neural net is defined as
actions * atoms. forwardDist(…)is responsible for generating the probability distributionp(s, a).forwardLogits(…)is used when you need to train, because the loss will be Negative Log Likelihood Loss21, and in the internal implementation, it must receive logits22 so it can apply log-softmax (for numerical stability).qValuesFromDist(…)is used together withforwardDist(…)to applyargMax(…)at the end to extract the action with the highest Q-value.- We call
projectBellman(…)to generate the target distribution. newAtomsBroadcaster(…)serves to generate anNDArraythat allows the code to produce the shape (batch, 1, atoms) so we can use thegather(…)method passing theactionindex.- The
forward(…)default method fromICategoricalNetworkalready usesforwardDist(…)andqValuesFromDist(…)together.
// omitted code ...
public class CategoricalQNetworkMLP implements IDeepQNetwork, ICategoricalNetwork {
private boolean training;
private final NDManager subManager;
private final int observations;
private final int actions;
private final int atoms;
private final CategoricalBellmanProjection catProj;
private final Model model;
private final SequentialBlock net;
private final ParameterStore parameterStore;
private final NDArray support;
// omitted code ...
@SneakyThrows
public CategoricalQNetworkMLP(int observations,
int actions,
CategoricalBellmanProjection catProj,
Path modelPath,
String prefix,
NDManager parent) {
this.observations = observations;
this.actions = actions;
this.catProj = catProj;
this.subManager = subMgr(parent, getClass());
// Configure C51 Parameters
this.atoms = catProj.atoms();
this.support = catProj.support(subManager);
this.model = newModel(getClass(), subManager.getDevice());
this.net = new SequentialBlock();
net.add(linear(128))
.add(Activation::relu)
.add(linear(128))
.add(Activation::relu)
.add(linear((long) actions * atoms));
model.setBlock(net);
// omitted code ...
}
@Override
public NDArray forwardDist(NDList inputs, final UnaryOperator<NDArray> block) {
@Cleanup var logits = safeForwardSingle(subManager, net, parameterStore, inputs, training).singletonOrThrow();
return scoped(it -> {
var probDist = it.reshape(N_BATCH, actions, atoms)
.softmax(LAST_AXIS);
var out = block.apply(probDist);
out.tempAttach(it.getManager());
return out;
}, logits);
}
@Override
public NDArray forwardLogits(NDArray input, final UnaryOperator<NDArray> block) {
@Cleanup var logits = safeForwardSingle(subManager, net, parameterStore, new NDList(input), training).singletonOrThrow();
return scoped(it -> {
var logitsDist = it.reshape(N_BATCH, actions, atoms);
var out = block.apply(logitsDist);
out.tempAttach(it.getManager());
return out;
}, logits);
}
@Override
public NDArray qValuesFromDist(NDArray distribution) {
if (distribution.getShape().dimension() != 3) {
throw new IllegalStateException("Invalid shape, must be (batch, actions, atoms)!");
}
// (batch, actions, atoms) * (1, atoms) -> (batch, actions)
return scoped(array -> {
var dist = array[0];
var sup = array[1];
return dist.mul(sup).sum(LAST_AXIS_ARR);
}, distribution, support);
}
@Override
public NDArray projectBellman(final NDArray probNextDist,
final NDArray rewards,
final NDArray dones,
final float gamma) {
return catProj.project(probNextDist, rewards, dones, gamma);
}
public int atoms() {
return atoms;
}
public NDArray newAtomsBroadcaster(NDManager external) {
return external.ones(new Shape(1, 1, atoms))
.stopGradient();
}
// omitted code ...
}
Categorical C51 Loss Implementation
The class CategoricalNLLLoss implements the loss function Negative Log Likelihood, as explained before, it’s more stable when you’re training networks, but it’s mandatory that the output from dist that comes from online-network must be logSoftmax(…) instead of softmax(…). The softmax(…) function is used to get the probabilities distribution, where the sum of all outputs is 1:
// omitted code ...
public class CategoricalNLLLoss extends Loss {
public CategoricalNLLLoss() {
super(CategoricalNLLLoss.class.getSimpleName());
}
@Override
public NDArray evaluate(NDList massDist, NDList distLogits) {
// (batch, 1, atoms)
var predLogits = distLogits.singletonOrThrow();
// (batch, 1, atoms)
var lab = massDist.singletonOrThrow();
// (batch, 1, 1) - pred must be log softmax
var loss = predLogits.logSoftmax(LAST_AXIS)
.mul(lab)
.neg()
.sum(LAST_AXIS_ARR, KEEP_DIMS);
// () -> scalar
return loss.mean();
}
}
Training Loop
Now we compose the AgentC51DQN:
- We call
trainStepClipGradients(…)to avoid gradient explosion.
// omitted code ...
public class AgentC51DQN extends AbstractAgent {
private static final float CLIP_GRAD_THRESHOLD = 10.0f;
private final CategoricalQNetworkMLP onlineCatNet;
private final CategoricalQNetworkMLP targetCatNet;
private final NDArray atomsBroadcaster;
private final NDManager subManager;
// omitted code ...
@Override
protected float trainOnline(int batchSize,
IReplayBuffer replayBuffer,
Loss lossFunc,
NDManager sub) {
if (!replayBuffer.enough(batchSize)) return Float.NaN;
if (!(lossFunc instanceof CategoricalNLLLoss)) {
throw new IllegalArgumentException("You must pass CategoricalNLLLoss!");
}
@Cleanup var samples = replayBuffer.sample(batchSize);
@Cleanup var projectDist = targetCatNet.forwardDist(samples.nextStates(), probNextDist -> {
// (batch, actions, atoms) -> (batch, actions)
var nextQValues = targetCatNet.qValuesFromDist(probNextDist);
var maxNextActions = nextQValues.argMax(AXIS_1)
// (batch, 1, 1)
.reshape(N_BATCH, 1, 1)
// (batch, 1, atoms)
.mul(atomsBroadcaster);
// (batch, 1, atoms) - now we are really selecting only actions a*, not all actions -> p(s’, a*, theta-).
var maxNextProbDist = probNextDist.gather(maxNextActions, AXIS_1).stopGradient();
// Bellman Projection - mi
return targetCatNet.projectBellman(maxNextProbDist, samples.rewards(), samples.dones(), gamma);
});
// Loss = -1/n sum mi * ln (p(s, a, theta))
float lossItem = backwardLoss(sub, lossFunc, projectDist, array -> {
var states = array[0];
var actions = array[1]
// (batch, 1, 1)
.reshape(N_BATCH, 1, 1)
// (batch, 1, atoms)
.mul(atomsBroadcaster);
// logits z(s, a, theta)
return onlineCatNet.forwardLogits(states, logits -> logits.gather(actions, AXIS_1));
}, samples.states(), samples.actions());
trainStepClipGradients(onlineNet.getBlock(), optimizer, CLIP_GRAD_THRESHOLD);
return lossItem;
}
// omitted code ...
}
Step-by-Step Value Transformation Explanation
Let’s see some values step-by-step:
Here we are applying the reward and the discount factor to each element in the support vector (each position is an atom). We do this because it’s a way to incorporate the values that we got from the environment and make sure that the network learns this distribution. Let’s imagine that we wouldn’t do that: we would not learn how the environment returns its rewards for each state. Furthermore, we limit the values between Vmin and Vmax to make the training more stable, as it prevents overflow and prevents extreme values from propagating through the network during backpropagation:
Atom Projection: Shift and Clip Each Atom
At the top of the table, you can see the distribution being applied to each atom. At first, it looks like we are just replicating the same value that increases because at each atom we have a bigger value. However, when the online network learns this distribution through backpropagation, and at fixed timesteps we update the target network with the online network’s weights, we ensure that both networks evolve over time. The target network does not backpropagate because we want to keep it stable; this avoids the “dog chasing its tail” effect, where if both networks update simultaneously, there is no stable reference, making convergence harder or even impossible:
Mass Redistribution via Linear Interpolation
Starting the mass redistribution, where the split happens between the atoms, between the lower (l) and upper (u), a linear interpolation23 happens between each atom. The case l = u is the edge case, where we just add pj (without multiplying by u or l):
Note the application of the linear interpolation formula A + (B - A) * t
Below is the proof that the sum of all contributions into m is equal to 1 (100% of the probability):
KL Divergence and Cross-Entropy in C51
The image below shows the Kullback-Leibler being applied between the m (Bellman projection from target network) and p(s, a) (from the online network), in practice, we use the Cross-Entropy to train the network, in this case, they behave the same way:
In the formula below, you can see that the “Negative Entropy of Target” part calculates the Entropy of the Bellman projection m. However, because we do not update the target network during backpropagation, we treat this entire term as a constant and drop it from the computational graph. In the end, the optimization depends only on the Cross-Entropy, which also makes training faster. You can see the two forms of the Kullback-Leibler formula in this link. Check the formula below:
Distributional RL vs. Standard DQN
Distributional RL (C51) maintains the return distribution, enforced by the model architecture and the way we train it; given enough time, the model will converge to the true return distribution. This differs from standard DQN, where the model has no architectural constraint: it may implicitly represent some distribution internally, but we do not guarantee this.
C51, on the other hand, uses fixed atoms zi on a support [Vmin, Vmax] with a column-wise softmax to structurally enforce a valid probability distribution over returns (for numerical stability, log-softmax is used with the log-sum-exp trick24). Check the example below to make it more clear:
We have below an environment where the treasure has a positive reward of +100, and a trap with a negative reward of -100.
When we calculate the Bellman equation, this distribution disappears into an average value:
But, different from C51, this value is maintained in the distribution, and the network is ensured to learn the distribution and respect it over time:
The expected value E[Z(s,a)] yields the same scalar as DQN’s Q-value, but is only computed at inference time for action selection. Backpropagation operates directly on the full categorical distribution via cross-entropy loss, which is why uncertainty, risk, and variance are fully preserved throughout training. Standard DQN collapses the entire return distribution into a single scalar, and even if some distributional structure emerges implicitly in the weights during training, there is no guarantee it will, and the output will always discard it.
Architecture Comparison: DQN, Dueling DQN, and C51
Now you can see the difference in architecture and output of DQN, Dueling DQN and C51 (note that in C51, for each action, it has 8 atoms in the image below, but in the standard C51, they used 51):
Difference between DQN, Dueling-DQN and C51 when selecting actions (without Noisy Networks):
Note that the output shapes of DQN and Dueling DQN will be (actions,), but C51’s will be (actions, 8), which collapses back to (actions,) once the expected value is computed.
Practical Tuning: Choosing Vmin and Vmax
I tested the value vMin = -50.0 and vMax = +50.0 and the AgentC51 didn’t learn effectively (points between -20 and at most 100); but when I put vMin = -100.0 and vMax = +100.0, it worked and the reward passed 200 points. The cause is that we were clipping the distribution too much, so the algorithm wasn’t learning the approximated distribution.
Rainbow DQN
The Rainbow DQN article25 implements all the algorithms and structure that we have learned though this article. As the other implementation was explained before, we’ll proceed to the formula then the code of Rainbow only.
Categorical C51 Loss Implementation with PER
The loss is the NLLLoss with the addition of PER Loss too:
In the code below we use logSoftmax(…) too:
// class CategoricalNLLPERLoss
// omitted code ...
@Override
public NDArray evaluate(NDList massDist, NDList distLogits) {
if (null == normISWeights) {
throw new IllegalStateException(“You must set the normISWeights before training!”);
}
// (batch, 1, atoms)
var predLogits = distLogits.singletonOrThrow();
// (batch, 1, atoms)
var lab = massDist.singletonOrThrow();
// (batch, 1, 1) - pred must be logits
var loss = normISWeights.mul(predLogits.logSoftmax(LAST_AXIS)
.mul(lab)
.sum(LAST_AXIS_ARR, KEEP_DIMS))
.neg();
this.normISWeights = null;
if (reduction.equals(Reduction.MEAN)) {
return loss.mean();
}
// (batch) - We need to extract the loss per element to update sum/min segment-tree
return loss.squeeze();
}
// omitted code ...
Rainbow Q-Network with Dueling and Noisy Layers
The RainbowQNetworkMLP is similar to the CategoricalQNetworkMLP, but using the DuelingCategoricalLayer with NoisyLayers instead of a SequentialBlock of standard linear layers:
DuelingCategoricalLayerimplements the Dueling DQN pattern but adapted toC51.- NoisyLayers comes from Noisy Networks, which we covered in this article.
// omitted code ...
public class RainbowQNetworkMLP implements INoisyNetwork, ICategoricalNetwork {
// omitted code ...
@Getter
private final DuelingType duelingType;
private final List<NoisyLayer> noisyLayers;
private final Model model;
private final DuelingLayer net;
// omitted code ...
@SneakyThrows
public RainbowQNetworkMLP(int observations,
int actions,
CategoricalBellmanProjection catProj,
Path modelPath,
String prefix,
NDManager parent,
DuelingType duelingType) {
// omitted code ...
this.net = new DuelingCategoricalLayer(
actions,
atoms,
duelingType,
featureBackbone ->
featureBackbone.add(linear(128))
.add(Activation::relu)
.add(addAndGet(noisyLayers, noisyLayer(128)))
.add(Activation::relu)
.add(addAndGet(noisyLayers, noisyLayer(128)))
.add(Activation::relu),
valueHead ->
valueHead.add(addAndGet(noisyLayers, noisyLayer(atoms))),
advantageHead ->
advantageHead.add(addAndGet(noisyLayers, noisyLayer(actions * atoms)))
);
model.setBlock(net);
// omitted code ...
}
@Override
public NDArray forwardDist(NDList inputs, final UnaryOperator<NDArray> block) // omitted code ...
@Override
public NDArray forwardLogits(NDArray input, final UnaryOperator<NDArray> block) // omitted code ...
@Override
public NDArray qValuesFromDist(NDArray distribution) // omitted code ...
@Override
public NDArray projectBellman(final NDArray probNextDist,
final NDArray rewards,
final NDArray dones,
final float gamma) // omitted code ...
}
Combining Dueling with C51
The class DuelingCategoricalLayer overrides the DuelingLayer but modifies the output shape to match (batch, actions, atoms):
// omitted code ...
public class DuelingCategoricalLayer extends DuelingLayer {
private final int atoms;
public DuelingCategoricalLayer(int actions,
int atoms,
DuelingType duelingType,
Consumer<SequentialBlock> featureBackboneModifier,
Consumer<SequentialBlock> valueHeadModifier,
Consumer<SequentialBlock> advantageHeadModifier) {
super(actions, duelingType, featureBackboneModifier, valueHeadModifier, advantageHeadModifier);
this.atoms = atoms;
}
@Override
public NDList forwardInternal(ParameterStore parameterStore, NDList inputs, boolean training, PairList<String, Object> params) {
if (inputs.size() != 1) throw new IllegalArgumentException(“The dueling categorical dqn just accepts one input!”);
var feature = featureBackbone.forward(parameterStore, inputs, training);
// (batch, 1, atoms)
var value = valueHead.forward(parameterStore, feature, training).singletonOrThrow()
.reshape(N_BATCH, 1, atoms);
// (batch, actions, atoms)
var advantage = advantageHead.forward(parameterStore, feature, training).singletonOrThrow()
.reshape(N_BATCH, actions, atoms);
// (batch, actions, atoms)
return new NDList(qValueCalc.apply(value, advantage));
}
// omitted code ...
}
Full Integration
The AgentRainbowDQN uses the NStepPrioritizedReplayBuffer that is a combination of PrioritizedReplayBuffer with the n-step learning algorithm:
- inside
trainOnline(…)you have DDQN, C51 and Noisy Networks (because ofresetNoise()).- The
onlineRainbowNet.forward(nextStates, …)is the DDQN part. targetRainbowNet.forwardDist(…)obtains the target network’s categorical distribution, whileprojectBellman(…)constructs the n-step Bellman target using the bootstrap term.catLossFuncis the NLLLoss with PER.
- The
selectAction(…)uses Noisy Networks to sample or select the best action effectively.templateExtraProcessing(…)is responsible for adjusting beta to increase Importance Sampling correction over time.
// omitted code ...
public class AgentRainbowDQN extends AbstractAgent {
private static final float CLIP_GRAD_THRESHOLD = 10.0f;
private final RainbowQNetworkMLP onlineRainbowNet;
private final RainbowQNetworkMLP targetRainbowNet;
private final NDManager subManager;
private final NDArray atomsBroadcaster;
private final float initialBeta;
private float beta;
// omitted code ...
@Override
protected float trainOnline(int batchSize, IReplayBuffer ireplayBuffer, Loss lossFunc, NDManager sub) {
if (!ireplayBuffer.enough(batchSize)) return Float.NaN;
if (!(ireplayBuffer instanceof NStepPrioritizedReplayBuffer replayBuffer)) {
throw new IllegalArgumentException(“You must pass NStepPrioritizedReplayBuffer!”);
}
if (!(lossFunc instanceof CategoricalNLLPERLoss catLossFunc)) {
throw new IllegalArgumentException(“You must pass CategoricalNLLPERLoss!”);
}
// PER
@Cleanup var samples = replayBuffer.sample(batchSize, beta);
// reset first time eps’ for DDQN
onlineRainbowNet.resetNoise();
// a* = arg max q_online(s’, a’)
var nextStates = samples.nextStates();
@Cleanup var action = onlineRainbowNet.forward(nextStates,
onlineNextQValues -> onlineNextQValues.stopGradient().argMax(AXIS_1)
// (batch, 1, 1)
.reshape(N_BATCH, 1, 1)
// (batch, 1, atoms)
.mul(atomsBroadcaster));
@Cleanup var targetQValue = targetRainbowNet.forwardDist(nextStates, probNextDist -> {
float gammaNBootstrap = (float) Math.pow(gamma, replayBuffer.nStep());
// q_target(s’, a*) - DDQN
var bestNextProbDist = probNextDist.gather(action, AXIS_1);
// Bellman Projection - mi (with n-step gammaNBootstrap instead of just gamma)
return targetRainbowNet.projectBellman(bestNextProbDist, samples.rewards(), samples.dones(), gammaNBootstrap);
});
// reset second time eps’ for DDQN now to compute TD-Error
onlineRainbowNet.resetNoise();
catLossFunc.normISWeights(samples.weights());
var losses = rawBackwardLoss(sub, catLossFunc, targetQValue, array -> {
var states = array[0];
var actions = array[1]
// (batch, 1, 1)
.reshape(N_BATCH, 1, 1)
// (batch, 1, atoms)
.mul(atomsBroadcaster);
// logits z(s, a, theta)
return onlineRainbowNet.forwardLogits(states, logits -> logits.gather(actions, AXIS_1));
}, samples.states(), samples.actions());
var lossItem = scopedToFloat(NDArray::mean, losses);;
@Cleanup var priorities = scoped(it -> it.abs().add(MIN_PRIORITY), losses);
replayBuffer.updatePriorities(samples.bufferIndexes(), priorities);
trainStepClipGradients(onlineRainbowNet.getBlock(), optimizer, CLIP_GRAD_THRESHOLD);
return lossItem;
}
@Override
public ActionSpaceType.ActionResult selectAction(NDArray state) {
// Noisy Networks
onlineRainbowNet.resetNoise();
@Cleanup var oneBatchState = state.expandDims(0);
long action = onlineRainbowNet.forwardLong(oneBatchState,
qValue -> qValue.stopGradient().argMax(AXIS_1));
return actionSpaceType.get(action);
}
@Override
protected void templateExtraProcessing(int frames, long frameLimit) {
// PER
float fraction = Math.min((float) frames / frameLimit, 1.0f);
this.beta = this.initialBeta + fraction * (1.0f - this.initialBeta);
}
}
Finally:
AgentRainbowDQN in the Lunar Lander environment:After that, we’ll use the library carla-simulator-javacpp to make the integration with the CARLA Simulator26 (the article was first posted on substack).
We conclude the implementation of the main algorithms associated with Rainbow DQN, note that we can add and do many more modifications, like using deeper networks or even transformer networks.
- Deep Java Library: https://github.com/deepjavalibrary/djl. ↩︎
- Lunar Lander: https://gymnasium.farama.org/environments/box2d/lunar_lander/. ↩︎
- Region-based memory management: https://en.wikipedia.org/wiki/Region-based_memory_management. ↩︎
- DJL Memory Management good practices: https://docs.djl.ai/master/docs/development/memory_management.html. ↩︎
- Article explaining about the DQN: https://neuralpalace.substack.com/p/how-to-never-forget-deep-q-networks ↩︎
- Article explaining about Q-Learning: https://neuralpalace.substack.com/p/video-games-and-reinforcement-learning. ↩︎
- Transfer learning: https://en.wikipedia.org/wiki/Transfer_learning. ↩︎
- “Deep Reinforcement Learning with Double Q-learning” by Hado van Hasselt et al. (Google DeepMind, 2015) – https://arxiv.org/abs/1509.06461. ↩︎
- “Dueling Network Architectures for Deep Reinforcement Learning” by Ziyu Wang et al. (Google DeepMind, 2016) – https://arxiv.org/abs/1511.06581. ↩︎
- Huber Loss: https://en.wikipedia.org/wiki/Huber_loss. ↩︎
- Mean Absolute Error: https://en.wikipedia.org/wiki/Mean_absolute_error. ↩︎
- “Prioritized Experience Replay” by Tom Schaul et al. (Google DeepMind, 2015) – https://arxiv.org/abs/1511.05952. ↩︎
- Monte Carlo method: https://en.wikipedia.org/wiki/Monte_Carlo_method. ↩︎
- “Rainbow: Combining Improvements in Deep Reinforcement Learning” by Matteo Hessel et al. (Google DeepMind, 2017) – https://arxiv.org/abs/1710.02298. ↩︎
- More about Normal Distribution: https://en.wikipedia.org/wiki/Normal_distribution. ↩︎
- Article Understanding the Reparameterization Trick: https://medium.com/@ml_dl_explained/understanding-the-reparameterization-trick-be349756b91b ↩︎
- Other usages: https://en.wikipedia.org/wiki/Reparameterization_trick. ↩︎
- “Noisy Networks for Exploration” by Meire Fortunato et al. (Google DeepMind, 2017) – https://arxiv.org/abs/1706.10295. ↩︎
- “A Distributional Perspective on Reinforcement Learning” by Marc G. Bellemare et al. (Google DeepMind, 2017) – https://arxiv.org/abs/1707.06887. ↩︎
- Article about Entropy, Cross-Entropy and Kullback-Leibler Divergence: https://www.datacamp.com/tutorial/kl-divergence. ↩︎
- Negative Log Likelihood Loss (NLLLoss) explanation: https://docs.pytorch.org/docs/2.13/generated/torch.nn.NLLLoss.html. ↩︎
- Logits are the raw (non-normalized) outputs of a neural network: https://stackoverflow.com/questions/41455101/what-is-the-meaning-of-the-word-logits-in-tensorflow ↩︎
- You can learn more about linear interpolation (lerp), in the following links: Explained: Linear Interpolation [Math] and An In-Depth look at Lerp, Smoothstep, and Shaping Functions. ↩︎
- Explanation of Log-Sum-Exp trick can be found on the following links:
https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/ and https://en.wikipedia.org/wiki/LogSumExp ↩︎ - “Rainbow: Combining Improvements in Deep Reinforcement Learning” by Matteo Hessel et al. (DeepMind, 2017) – https://arxiv.org/abs/1710.02298. ↩︎
- CARLA Simulator version 0.96.0: https://carla.org/ and https://github.com/carla-simulator/carla/releases/tag/0.9.16 ↩︎
