Skip to content

Pqn

PQN ("Parallelised Q-Network"): an online, parallel-environment deep Q-learning agent with no replay buffer and no target network.

Standard DQN needs both a replay buffer (to decorrelate updates from a single, highly autocorrelated trajectory) and a target network (to keep the regression target from chasing itself as the online network updates) to avoid diverging. PQN's core claim is that neither is necessary once the Q-network is regularized with LayerNorm and trained on data from many parallel environments at once (as PPO already does for the same decorrelation reason): the online network's own current weights can serve as the bootstrap target, and LayerNorm keeps that self-referential regression stable instead of blowing up. What survives is close to the simplest thing that could be called deep Q-learning: collect a rollout, regress Q(s, a) towards a bootstrapped return, repeat.

Per-update loop (update): 1. collect_experience - hparams.num_steps steps across hparams.num_envs parallel envs, epsilon-greedy over the online network's own Q-values (distrax.EpsilonGreedy, annealed via epsilon). Each transition caches max_a Q(s_t, a) (Buffer.value) at collection time - the reference implementation's values array - so the target computation below never needs to re-run the network over already-visited states. 2. evaluate_experience - a Q(lambda) return over the whole rollout, mixing the one-step bootstrap ((1 - q_lambda), using the cached value one step ahead) with the multi-step return (q_lambda, chaining through the target, not the cached value, at one step ahead) - literally rlax.lambda_returns, whose own docstring calls out this exact use: "Q(lambda): v_t = max(q_t, axis=-1)". Computed once per rollout, not re-evaluated per epoch the way PPO.update re-evaluates its GAE targets - the reference implementation runs update_epochs of SGD against one fixed set of targets, since there's no importance-ratio correction here that would make re-evaluating them mid-epoch meaningful. 3. hparams.num_epochs passes of shuffled-minibatch MSE regression (q_loss) of Q(s_t, a_t) towards that fixed target.

No target network: the bootstrap value at every step - both the cached Buffer.value and the final bootstrap in evaluate_experience - comes from the same train_state.params the rollout was collected with, not a separately-updated copy. No replay buffer: every minibatch this update draws on is a shuffled slice of the rollout collect_experience just produced, used once, then discarded - unlike DQN, nothing here persists across update calls.

On PQNHparams' defaults for gridworld tasks: num_epochs/ num_minibatches/exploration_fraction/end_e are set higher/lower than CleanRL's CartPole-tuned reference script, not because that script is wrong, but because a budget-frame run buys many fewer rollouts here (num_updates = budget // (num_steps * num_envs)) than CartPole's own total_timesteps example uses, and PQN gets no benefit from the extra fixed-target minibatch passes a replay-buffer method would - the only way to extract more learning signal per rollout is more epochs over it. Verified empirically (not just reasoned about): the external rejax package's own PQN, with its per-environment-tuned config for Navix-Empty-6x6-v0 (num_epochs=8, num_minibatches=128, exploration_fraction=0.3, end_e=0.1, gamma=0.9), reaches ~100% success where this file's original CleanRL-derived defaults reached ~72% at 1M frames on the (comparably simple) Navix-Empty-5x5-v0 - and dropping rejax's exact config into this implementation reproduces its ~100% result, confirming the target computation itself was never the problem. rejax's own target computation, on inspection (rejax/algos/pqn.py), turned out to diverge from the official reference it's adapted from (Gallici's own mttga/purejaxql/purejaxql/pqn_gymnax.py): the official version caches Q at the state a transition starts from and shifts it forward by one step when bootstrapping (this module's Buffer.value does the same); rejax's adaptation caches Q at the state the transition lands in instead, which is the wrong operand one step early in the backward recursion. This module's evaluate_experience follows the official convention, not rejax's.

Bases: PyTreeNode

The encoder carry (navix.agents.models.Encoder) that went into the network at this step - the pre-step state, as in navix.agents.ppo.Buffer.carry. () for the stateless Q-encoders; for TransformerEncoder it's the raw frame window, whose replay in q_loss is exact (parameter-independent).

Bases: Agent

Linear anneal from start_e to end_e over exploration_fraction * budget frames, then held at end_e.

The Q(lambda) return target, computed once per rollout (not re-evaluated per epoch - see this module's docstring). values are the cached max_a Q(o_t, a) from collection time (Buffer.value); the one bootstrap not already cached is max_a Q(o_T, a) at the post-rollout observation, under the same (not-yet-updated-this-round) params.

Bases: HParams

Hyperparameters for PQN. Frozen; .replace(...) for a variant. Several defaults (num_epochs, num_minibatches, end_e, exploration_fraction) are set for navix's gridworlds and differ from CleanRL's CartPole reference - see this module's docstring.

Whether to anneal the learning rate linearly to 0 at the end of training.

Number of environment frames to train for.

Final epsilon for epsilon-greedy exploration. Higher than CleanRL's CartPole-tuned default (0.05) - see this module's docstring.

Fraction of budget over which epsilon anneals from start_e to end_e; held at end_e for the remainder. Shorter than CleanRL's CartPole-tuned default (0.5) - see this module's docstring.

Hidden layer size of the Q-network.

Starting learning rate.

Maximum norm for gradient clipping.

Number of parallel environments to run.

Number of shuffled-minibatch passes per update over the rollout's (fixed, not re-evaluated) Q(lambda) targets - "update_epochs" in the reference implementation. Higher than CleanRL's CartPole-tuned default (4) - see this module's docstring.

Number of minibatches to split the rollout into. Higher than CleanRL's CartPole-tuned default (8) - see this module's docstring on gridworld-appropriate defaults for why.

Number of steps to run in each environment per update.

Mixing parameter for the Q(lambda) return target - see rlax.lambda_returns.

Initial epsilon for epsilon-greedy exploration.

Bases: TrainState

The live encoder carry for the num_envs running envs, threaded across collect_experience like env_state; also the pre-step carry for the post-rollout bootstrap. () for the stateless Q-encoders.