Jormungandr: Distributed C51 Reinforcement Learning

Jormungandr 0.1.0 is an asynchronous actor–learner service for discrete-action, off-policy reinforcement learning. Environment processes act independently and submit their experience to a central Python runtime. That runtime owns replay, parameter updates, held-out evaluation, policy versions, checkpoints, and inference.

In this release, C51 is the only algorithm connected to the complete service lifecycle. PPO and DDPG implementations are included as lower-level research code, but they do not yet have the rollout, replay, and concurrency contracts provided for C51. The repository does not count the presence of an algorithm class as operational support.

Jormungandr is released under Apache-2.0. The Python distribution is jormungandr-rl; its import name is jormungandr.

Actor–learner organization

Let an actor observe \(o_t\), choose action \(a_t\) under policy \(\pi_v\), receive reward \(r_t\), and observe \(o_{t+1}\). Jormungandr represents the resulting transition as

\[\tau_t = (o_t, a_t, r_t, o_{t+1}, d_t; \text{actor}, \text{episode}, t, \text{split}, v),\]

where \(d_t\) records termination and \(v\) is the policy version used to choose the action. The fields after the semicolon are not annotations added during analysis. They are required by the public ingestion protocol.

The actor retains control of episode construction, environment state, and reward calculation. It obtains an action through the inference endpoint, steps its environment, and submits the transition to /v1/models/{model_id}/experience/add. Inference returns both the action index and the corresponding domain value, together with expected action values and the current policy version.

Versioning is necessary because acting and learning proceed concurrently. While an actor is completing an episode with \(\pi_v\), the learner may advance to \(\pi_{v+1}\) or beyond. Storing \(v\) makes this policy lag observable. Version 0.1.0 does not yet reject stale transitions or apply a correction for policy lag; it records the information required to study or impose such a rule later. Nor is a version number a behavior-policy likelihood. The current experience record does not contain the action probability required for a general importance-ratio correction between actor policies.

Several actors may contribute to one request, and training and validation items may be interleaved. Routing is therefore a property of each transition, not of the connection through which it arrived. The experience protocol requires actor_id, episode_id, timestep, split, and policy_version. The same identity is used to attach an auxiliary label that arrives before or after its transition.

This organization belongs to the general family of distributed replay systems in which acting is decoupled from learning. Jormungandr's present implementation is deliberately smaller than a system such as Ape-X: replay is process-local, and each model has one authoritative learner rather than a synchronized population of learners.

Replay and holdout

Training transitions enter a prioritized replay buffer. If \(p_i\) is the current priority of item \(i\), sampling follows

\[P(i) = \frac{p_i^\alpha} {\sum_j p_j^\alpha}.\]

The loss for a sampled transition is multiplied by a normalized importance-sampling weight

\[w_i = \frac{(N P(i))^{-\beta}} {\max_j (N P(j))^{-\beta}},\]

where \(N\) is the replay size and \(\beta\) increases towards one during training. This is the method introduced by Schaul et al.. In Jormungandr, the updated priority is derived from the per-transition categorical loss.

Validation transitions are placed in a separate bounded store. They do not enter prioritized replay, alter priorities, contribute gradients, or update the running observation mean and variance. They are transformed using the normalization statistics estimated from training data.

The repository tests this rule with one mixed batch:

train -> validation -> train -> validation

The resulting replay and validation stores each contain two transitions. The normalizer count is four because each training transition contributes \(o_t\) and \(o_{t+1}\); the validation observations contribute nothing. A separate test waits for a learner update and verifies that the held-out loss is reported with the policy version under which it was evaluated. These assertions appear in the runtime tests.

The test concerns sample isolation alone. It establishes that a held-out sample cannot change the fitted normalization or the model before its loss is reported; it provides no evidence that a particular validation design has enough observations or covers the relevant environment distribution.

There are also two held-out quantities in the OU example. The learner reports categorical loss on transitions in the validation store, measuring agreement with its projected one-step Bellman targets for the recorded actions. The example separately totals reward over deterministic, closed-loop validation episodes. One is a transition-level training diagnostic; the other concerns the behavior of the fitted policy over an episode.

Categorical value learning

Conventional Q-learning estimates the expected discounted return

\[Q(s,a) = \mathbb{E}\left[ \sum_{k=0}^{\infty}\gamma^k R_{t+k+1} \mid S_t=s, A_t=a \right].\]

C51 instead approximates the distribution of that return. In its standard form, the network assigns probabilities to 51 equally spaced atoms

\[z_j = V_{\min} + j\Delta z, \qquad \Delta z = \frac{V_{\max}-V_{\min}}{50}, \qquad j=0,\ldots,50.\]

If \(p_j(s,a)\) is the probability assigned to atom \(z_j\), the action value used for selection remains the expectation

\[Q(s,a) = \sum_{j=0}^{50} z_j p_j(s,a).\]

Jormungandr exposes the atom count as a configuration parameter; both its default and the OU experiment use 51. Although the model estimates a return distribution, action selection in this implementation maximizes the expectation above. It is not a risk-sensitive decision rule.

The distinction is in the training target. The distribution associated with the greedy next action is shifted to \(r_t+\gamma z_j\), or to \(r_t\) at termination, clipped to \([V_{\min},V_{\max}]\), and projected back onto adjacent atoms. The network is trained by cross-entropy against that projected distribution. This is the categorical algorithm introduced by Bellemare, Dabney, and Munos (2017).

Jormungandr implements the model as two ReLU hidden layers followed by number_of_actions × 51 logits. A separate target network supplies the projected distribution and is copied from the fitted network at a configured interval. Gradient clipping, prioritized replay, and an optional auxiliary classification loss are applied in the same update. Each completed update increments the policy version.

The default service support is \([-10,10]\). The limits are modelling choices: probability mass outside the interval is placed at an endpoint. A support that is too narrow therefore changes the target, while one that is unnecessarily wide reduces resolution between adjacent atoms. The OU example uses \([-8,8]\), \(\gamma=0.98\), a 64-unit hidden width, and a target-network update every 60 learner steps.

Ornstein–Uhlenbeck control experiment

The included experiment uses a discrete Ornstein–Uhlenbeck process as a synthetic mean-reverting spread:

\[x_{t+1} = (1-\kappa)x_t+\sigma\varepsilon_{t+1}, \qquad \varepsilon_{t+1}\sim\mathcal{N}(0,1).\]

With \(\kappa=0.18\) and \(\sigma=0.35\), its stationary standard deviation is

\[\sigma_x = \frac{\sigma} {\sqrt{1-(1-\kappa)^2}}.\]

For the stated parameters, \(\sigma_x\) is approximately \(0.612\).

The observation contains four quantities: \(x_t/\sigma_x\), the latest spread change in the same units, the current position, and the fraction of the 48-step episode remaining. The action set is \(\{-1,0,+1\}\), labelled SHORT, FLAT, and LONG. The chosen action is the position held over the next change in the spread.

Position is included because the turnover term depends on the previous action. Remaining time is included because the problem terminates after a fixed horizon. Omitting either variable would leave information relevant to the reward or transition law outside the observation presented to the learner.

The one-step reward is

\[r_t = a_t\frac{x_{t+1}-x_t}{\sigma_x} -0.035|a_t-a_{t-1}| -0.004|a_t|.\]

The second term charges changes in position; the third applies a small position penalty. Both are expressed in the same synthetic reward units as the standardized spread change. No estimate of market impact, bid–ask spread, margin, financing, or option lifecycle enters this environment.

A fixed reference policy is evaluated on each realized path. It is short above \(+0.65\sigma_x\), long below \(-0.65\sigma_x\), and flat between those levels. Because the learner and reference see the same random path, their episode rewards can be compared without path variation entering the difference.

Training episodes use epsilon-greedy action selection. Validation episodes use deterministic inference and independently seeded paths. Four training episodes are followed by one validation episode in the declared schedule, although the two actor threads may complete them in a different order.

The auxiliary classifier predicts the sign of conditional drift: DOWN when the spread z-score is above \(0.35\), UP below \(-0.35\), and FLAT between them. This target is closely related to a variable already present in the observation. High auxiliary accuracy consequently verifies the delayed-label and classification path; it should not be interpreted as independent forecast evidence.

Interpretation of the recorded run

The committed recording contains 28 training episodes and seven validation episodes, each of length 48. It therefore produces 1,344 training transitions and 336 validation transitions. The final frame reports 344 learner updates and 31 evaluations of the accumulated validation store.

The deterministic policy probe selects LONG at \(-2\sigma_x\) and \(-1\sigma_x\), FLAT at zero, and SHORT at \(+1\sigma_x\) and \(+2\sigma_x\). This is the ordering expected for control of a mean-reverting state. Mean reward across the seven validation episodes is 4.285, compared with 4.230 for the fixed reference on the same paths.

There are too few held-out episodes to attach much inferential weight to that difference, and the repository reports no confidence interval for it. Both policies are evaluated within the same known OU family used to construct the experiment. The result therefore shows that the distributed ingestion, learner, validation, and inference paths operate together; it does not show that C51 is superior to the reference, that the policy would generalize to a different stochastic process, or that it has value on market data.

The reported auxiliary accuracies—0.984 for training and 0.992 for validation—should be read in the same way. They are useful end-to-end checks of auxiliary training and held-out evaluation. The class boundary is, however, a threshold on the observed spread z-score, so the classification problem is intentionally simple.

Seeds and split assignments are fixed. Exact update counts may still vary slightly because experience ingestion and learner updates run concurrently. The experiment, reward definition, and recording procedure are documented in the OU spread note and reproducible results.

Scope of the release

Jormungandr supplies neither the Markov decision process nor its financial interpretation. Environments, reward functions, market data, portfolio rules, and execution remain actor or application responsibilities. This permits a Volt-based actor to use Jormungandr without placing option-contract semantics inside the learner.

Models have separate learner threads and locks, but replay and validation stores remain in memory within one process. They are not included in the 0.1 checkpoint format. A checkpoint retains network parameters, optimizer state, the training normalizer, feature order, update count, and policy version; it cannot by itself reproduce the exact replay state of an interrupted run.

The bundled HTTP service assumes a trusted network. Authentication, transport security, request admission, and rate limiting must be supplied by the deployment environment. These limits are set out in the architecture documentation.

The implementation can be installed and tested with:

python -m pip install -e ".[dev]"
python -m pytest

The current test suite contains 13 tests covering categorical projection, priority handling, split isolation, inference, checkpoints, exported TorchScript bundles, the generic trainer, and the OU environment.

References

• Marc G. Bellemare, Will Dabney, and Rémi Munos, A Distributional Perspective on Reinforcement Learning, Proceedings of Machine Learning Research 70, 2017.
• Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver, Prioritized Experience Replay, 2015.
• Dan Horgan et al., Distributed Prioritized Experience Replay, 2018.
• G. E. Uhlenbeck and L. S. Ornstein, On the Theory of the Brownian Motion, Physical Review 36, 1930.

Continue with Volt, the open-source options portfolio library, Hypercube, the open-source live analytics engine, or browse the Jormungandr source.