DQN (Deep Q-Network)¶

DQN is the classic reinforcement learning method approximating the Q-function with a neural network. The implementation uses a target network for stability and prioritized experience replay (PER) for more informative updates.

Components¶

Main Q-network: estimates \(Q_\theta(s,a)\) and selects actions
Target network: \(Q_{\theta^-}\) provides target values, updated less frequently
Replay buffer (PER): stores transitions and returns prioritized mini-batches
Action selection: \(\epsilon\)-greedy strategy

Theory¶

1) Bellman Optimality Equation¶

The optimal Q-function satisfies:

\[ Q^*(s,a) = \mathbb{E}\Big[r + \gamma \max_{a'} Q^*(s', a')\,\Big|\, s,a\Big] \]

Stochastic gradient descent on an MSE objective solves:

\[ \min_\theta \;\mathbb{E}_{(s,a,r,s')}\big[\big(y - Q_\theta(s,a)\big)^2\big],\quad \text{where } y = r + \gamma\, \max_{a'} Q_{\theta^-}(s', a') \]

2) Double DQN vs classical DQN¶

Vanilla DQN overestimates because the \(\max\) uses the same network. Double DQN decouples selection and evaluation:

\[ y = r + \gamma\, Q_{\theta^-}\Big(s', \operatorname*{argmax}\limits_{a'} Q_{\theta}(s', a')\Big) \]

This reduces overestimation and stabilizes training.

3) Target network update¶

The target network \(Q_{\theta^-}\) copies the online network every target_update_iter steps:

\[ \theta^- \leftarrow \theta \quad \text{(periodically)} \]

A fixed target over a short horizon reduces target drift.

4) Prioritized experience replay (PER) and SumTree¶

Transition priority i:

\[ p_i = |\delta_i| + \varepsilon_{\text{margin}} \quad \text{(may be clipped from above: } p_i \le \text{abs\_error\_upper)} \]

Sampling probability:

\[ P(i) = \frac{p_i^{\alpha}}{\sum_j p_j^{\alpha}}, \quad \alpha \in [0,1] \]

Importance-sampling weights:

\[ w_i = \Big( \frac{1}{N\, P(i)} \Big)^{\beta}, \quad \tilde{w}_i = \frac{w_i}{\max_j w_j}, \quad \beta \nearrow 1 \]

Priority update after training: \(p_i \leftarrow |\delta_i| + \varepsilon_{\text{margin}}\)
The SumTree structure enables \(\mathcal{O}(\log N)\) priority updates/sampling.

5) \(\epsilon\)-greedy policy and schedule¶

With probability \(\epsilon\) pick a random action; otherwise \(\arg\max_a Q_\theta(s,a)\).
Exploration decays: \(\epsilon \leftarrow \max(\text{min\_epsilon}, \epsilon \cdot \text{epsilon\_decay})\).

6) Training loop¶

Collect experience (first K steps without training).
Every replay_period steps sample a PER mini-batch, compute \(y\), TD errors \(\delta\), IS weights, and update \(\theta\) via weighted MSE.
Update priorities \(p_i\) and the parameter \(\beta\).
Every target_update_iter steps: \(\theta^- \leftarrow \theta\).

Pseudocode:

predict_q = Q_theta(s_batch)
best_action = argmax_a predict_q
target_q = Q_theta_minus(s_next_batch)
y = r_batch + gamma * target_q[range, best_action]

# TD error and priorities
delta = y - predict_q[range, a_batch]
priority = clip(|delta| + margin, 0, abs_error_upper) ** alpha

# importance weights and weighted MSE
w = ((buffer_size * P(i)) ** -beta) / max_w
loss = mean(w * (y - Q_theta(s_batch, a_batch))^2)
update theta by SGD

# update priorities, increase beta, periodically update target

7) Stabilization tricks¶

Gradient normalization/clipping
Limit TD error (abs_error_upper in code)
Prefer Huber loss over MSE (here we use weighted MSE)
Regular target updates, sufficiently large buffer

8) Mapping to implementation parameters¶

alpha — prioritization exponent (0 → uniform, 1 → pure TD error)
beta, beta_increment_per_sample — IS weight strength and annealing
target_update_iter — target network sync period
replay_period — training frequency
epsilon, epsilon_decay, min_epsilon — \(\epsilon\) schedule
margin (\varepsilon), abs_error_upper — priority shaping

Quick start¶

import gymnasium as gym
import numpy as np

from tensoraerospace.agent.dqn.model import Model, DQNAgent

env = gym.make('LinearLongitudinalF16-v0', number_time_steps=2000)
num_actions = env.action_space.n

model = Model(num_actions)
target_model = Model(num_actions)

agent = DQNAgent(
    model=model,
    target_model=target_model,
    env=env,
    train_nums=10000,
    epsilon=1.0,
    epsilon_dacay=0.995,
    min_epsilon=0.05,
)

agent.train()

Tip

For continuous action spaces use discretization or switch to DDPG/SAC.

Unified training interface¶

DQN supports the shared unified train() signature from BaseRLModel:

# Legacy call (still works):
agent.train()

# Unified call – overrides self.train_nums to num_episodes * max_steps:
agent.train(num_episodes=100, max_steps=200)

When num_episodes and max_steps are both omitted, the agent uses the train_nums step budget set at construction time.

API reference¶

`Model(num_actions)` ¶

Bases: Module

DQN with two hidden layers of 32.

Only the number of actions is required to preserve the original signature. The first linear layer is lazily initialized to infer the input features at runtime from the first forward pass.

Parameters:

Name	Type	Description	Default
`num_actions`	`int`	Number of actions.	required

Initialize network layers.

Parameters:

Name	Type	Description	Default
`num_actions`	`int`	Number of discrete actions.	required

`forward(x)` ¶

Compute Q-values for a batch of observations.

`predict(inputs)` ¶

Forward function. Returns Q-values for actions.

Parameters:

Name	Type	Description	Default
`inputs`	`ndarray`	Batch of input data (numpy array [B, obs_dim]).	required

Returns:

Type	Description
`ndarray`	numpy array [B, num_actions] with Q-values.

`action_value(obs)` ¶

Select greedy action and return Q-values for the first item.

Parameters:

Name	Type	Description	Default
`obs`	`ndarray`	Batch of input data.	required

Returns:

Name	Type	Description
`best_action`	`Union[ndarray, int]`	Best action(s). If batch size is 1 -> int, else ndarray.
`q_values`	`ndarray`	Q-values of the first element in the batch for compatibility.

`SumTree(capacity)` ¶

Binary search tree class for prioritized replay buffer agent.

Parameters:

Name	Type	Description	Default
`capacity`	`int`	Buffer size.	required

Initialize sum tree for prioritized replay.

Parameters:

Name	Type	Description	Default
`capacity`	`int`	Maximum number of transitions to store.	required

`total_p` `property` ¶

Number of records in buffer.

Returns:

Type	Description
`int`	Number of records in buffer.

`add(priority, transition)` ¶

Function for adding object to buffer.

Parameters:

Name	Type	Description	Default
`priority`	`int`	Priority of added transition.	required
`transition`	`Any`	Transition vector S, A, R, S'.	required

`update(idx, priority)` ¶

Function for updating object priority with given index.

Parameters:

Name	Type	Description	Default
`idx`	`int`	Transition index.	required
`priority`	`int`	Priority of updated transition.	required

`get_leaf(s)` ¶

Function for getting object by given priority.

Parameters:

Name	Type	Description	Default
`s`	`int`	Priority by which transition is selected.	required

Returns:

Name	Type	Description
`idx`	`int`	Transition index.
`priority`	`int`	Priority of updated transition.
`transitions`	`Any`	Required transition.

`DQNAgent(model, target_model, env, learning_rate=0.0012, epsilon=0.1, epsilon_dacay=0.995, min_epsilon=0.01, gamma=0.9, batch_size=8, target_update_iter=400, train_nums=5000, buffer_size=200, replay_period=20, alpha=0.4, beta=0.4, beta_increment_per_sample=0.001, log_dir=None, verbose_histogram=False, seed=1, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None)` ¶

DQN Agent.

Parameters:

Name	Type	Description	Default
`model`	`Module`	Deep Q-network model.	required
`target_model`	`Module`	Target deep Q-network model.	required
`env`	`Env`	Gym/Gymnasium environment.	required
`learning_rate`	`float`	Learning rate.	`0.0012`
`epsilon`	`float`	Environment exploration probability.	`0.1`
`epsilon_dacay`	`float`	Epsilon reduction coefficient per episode.	`0.995`
`min_epsilon`	`float`	Minimum epsilon value.	`0.01`
`gamma`	`float`	Discount coefficient.	`0.9`
`batch_size`	`int`	Mini-batch size.	`8`
`target_update_iter`	`int`	Target network update period (steps).	`400`
`train_nums`	`int`	Number of training steps.	`5000`
`buffer_size`	`int`	Replay buffer size.	`200`
`replay_period`	`int`	Buffer sampling period.	`20`
`alpha`	`float`	Prioritization degree.	`0.4`
`beta`	`float`	Importance sampling coefficient.	`0.4`
`beta_increment_per_sample`	`float`	Beta increment per sample.	`0.001`

Initialize DQN agent and replay buffer.

Parameters:

Name	Type	Description	Default
`model`	`Any`	Online Q-network.	required
`target_model`	`Any`	Target Q-network.	required
`env`	`Any`	Gym/Gymnasium environment.	required
`learning_rate`	`float`	Optimizer learning rate.	`0.0012`
`epsilon`	`float`	Initial epsilon for exploration.	`0.1`
`epsilon_dacay`	`float`	Multiplicative epsilon decay.	`0.995`
`min_epsilon`	`float`	Minimum epsilon value.	`0.01`
`gamma`	`float`	Discount factor.	`0.9`
`batch_size`	`int`	Training batch size.	`8`
`target_update_iter`	`int`	Steps between target updates.	`400`
`train_nums`	`int`	Total training steps to run.	`5000`
`buffer_size`	`int`	Replay buffer capacity.	`200`
`replay_period`	`int`	Sampling period from buffer.	`20`
`alpha`	`float`	PER priority exponent.	`0.4`
`beta`	`float`	PER importance sampling exponent.	`0.4`
`beta_increment_per_sample`	`float`	Increment for beta per sample.	`0.001`
`log_dir`	`str \| None`	Directory for TensorBoard logs.	`None`
`verbose_histogram`	`bool`	Whether to log histograms extensively.	`False`
`seed`	`int`	Random seed for numpy and torch. Defaults to 1.	`1`

`train(num_episodes=None, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)` ¶

Train the DQN agent (unified interface).

DQN is frame-based internally. num_episodes is treated as an episode target that is multiplied by max_steps to derive a total training-step budget. When max_steps is omitted, the original self.train_nums budget set at construction is used.

Parameters:

Name	Type	Description	Default
`num_episodes`	`Optional[int]`	Target number of episodes. Converted to a step budget via `num_episodes * max_steps` when `max_steps` is provided.	`None`
`max_steps`	`Optional[int]`	Approximate maximum steps per episode.	`None`
`save_best`	`bool`	Reserved for API consistency.	`False`
`save_path`	`Optional[str]`	Reserved for API consistency.	`None`
`verbose`	`bool`	Reserved for symmetry.	`True`
`**kwargs`	`Any`	Currently unused.	`{}`

Returns:

Name	Type	Description
`dict`	`dict`	Empty dict for API compatibility.

`train_step()` ¶

Function for training step.

Returns:

Name	Type	Description
`losses`	`float`	Losses after one training step.

`sum_tree_sample(k)` ¶

Get batch for training.

Parameters:

Name	Type	Description	Default
`k`	`int`	Size of batch to get.	required

Returns:

Name	Type	Description
`idxes`	`int`	Indices of objects from batch.
`is_weights`	`float`	Priorities of objects from batch.

`evaluation(wrapped_env, render=False)` ¶

Get batch for training.

Parameters:

Name	Type	Description	Default
`wrapped_env`	`Any`	Wrapped environment (for rendering/frame capture).	required
`render`	`bool`	Whether to visualize environment or not.	`False`

Returns:

Name	Type	Description
`ep_reward`	`float`	Total reward per episode.

`close()` ¶

Close resources (e.g., SummaryWriter).

`store_transition(priority, obs, action, reward, next_state, done)` ¶

Store a transition in the replay buffer.

`rand_based_sample(k)` ¶

Placeholder for rank-based prioritized sampling (not implemented).

`get_action(best_action)` ¶

Epsilon-greedy action selection.

`update_target_model()` ¶

Target neural network update function.

`get_target_value(obs)` ¶

Compute Q-values using the target network.

`e_decay()` ¶

Function for reducing network exploration probability.

`save(path=None, save_gradients=False)` ¶

Save PyTorch models to the specified directory.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path \| None`	Save path. If None, saves into current cwd under folder "dqn_agent".	`None`
`save_gradients`	`bool`	Save optimizer states to continue training.	`False`

`load(path, env, *, load_gradients=False, model_factory=None)` `classmethod` ¶

Load a DQNAgent from a directory created by :meth:save.

Parameters:

Name	Type	Description	Default
`path`	`Union[str, Path]`	Directory that contains `model.pth`, `target_model.pth` and `config.json` (as written by :meth:`save`).	required
`env`	`Any`	A Gymnasium-compatible environment. The environment is required because it is not serialised alongside the weights.	required
`load_gradients`	`bool`	If `True` and `optimizer.pth` exists, also restore the optimiser state so training can be continued.	`False`
`model_factory`	`Optional[Callable[[], Any]]`	Optional zero-argument factory for custom DQN network classes. Built-in `Model` checkpoints do not need it.	`None`

Returns:

Name	Type	Description
`DQNAgent`	`DQNAgent`	Fully initialised agent with loaded weights.

`from_pretrained(repo_name, env, access_token=None, version=None, load_gradients=False, model_factory=None)` `classmethod` ¶

Load pretrained model from local directory or Hugging Face Hub.

Parameters:

Name	Type	Description	Default
`repo_name`	`str`	Path to local folder with weights or repository name in format `namespace/repo_name` on Hugging Face Hub.	required
`env`	`Any`	A Gymnasium-compatible environment (required for reconstruction).	required
`access_token`	`Optional[str]`	Access token for private HF repository.	`None`
`version`	`Optional[str]`	Revision / branch / tag of HF repository.	`None`
`load_gradients`	`bool`	If `True`, also load optimiser states for continuing training.	`False`
`model_factory`	`Optional[Callable[[], Any]]`	Optional zero-argument factory for custom DQN network classes.	`None`

Returns:

Name	Type	Description
`DQNAgent`	`DQNAgent`	Initialised agent with loaded weights.

`publish_to_hub(repo_name, folder_path, access_token=None)` ¶

Publish saved model to Hugging Face Hub.

Parameters:

Name	Type	Description	Default
`repo_name`	`str`	Repository id on Hugging Face Hub (e.g. `"my-org/dqn-cartpole"`).	required
`folder_path`	`Union[str, Path]`	Local folder that contains saved weights (as written by :meth:`save`).	required
`access_token`	`Optional[str]`	HF access token. If `None`, the token cached by `huggingface-cli login` will be used.	`None`

DQN (Deep Q-Network)¶

Components¶

Theory¶

1) Bellman Optimality Equation¶

2) Double DQN vs classical DQN¶

3) Target network update¶

4) Prioritized experience replay (PER) and SumTree¶

5) \(\epsilon\)-greedy policy and schedule¶

6) Training loop¶

7) Stabilization tricks¶

8) Mapping to implementation parameters¶

Quick start¶

Unified training interface¶

API reference¶

Model(num_actions) ¶

forward(x) ¶

predict(inputs) ¶

action_value(obs) ¶

SumTree(capacity) ¶

total_p property ¶

add(priority, transition) ¶

update(idx, priority) ¶

get_leaf(s) ¶

train(num_episodes=None, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs) ¶

train_step() ¶

sum_tree_sample(k) ¶

evaluation(wrapped_env, render=False) ¶

close() ¶

store_transition(priority, obs, action, reward, next_state, done) ¶

rand_based_sample(k) ¶

get_action(best_action) ¶

update_target_model() ¶

get_target_value(obs) ¶

e_decay() ¶

save(path=None, save_gradients=False) ¶

load(path, env, *, load_gradients=False, model_factory=None) classmethod ¶

from_pretrained(repo_name, env, access_token=None, version=None, load_gradients=False, model_factory=None) classmethod ¶

publish_to_hub(repo_name, folder_path, access_token=None) ¶

`Model(num_actions)` ¶

`forward(x)` ¶

`predict(inputs)` ¶

`action_value(obs)` ¶

`SumTree(capacity)` ¶

`total_p` `property` ¶

`add(priority, transition)` ¶

`update(idx, priority)` ¶

`get_leaf(s)` ¶

`train(num_episodes=None, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)` ¶

`train_step()` ¶

`sum_tree_sample(k)` ¶

`evaluation(wrapped_env, render=False)` ¶

`close()` ¶

`store_transition(priority, obs, action, reward, next_state, done)` ¶

`rand_based_sample(k)` ¶

`get_action(best_action)` ¶

`update_target_model()` ¶

`get_target_value(obs)` ¶

`e_decay()` ¶

`save(path=None, save_gradients=False)` ¶

`load(path, env, *, load_gradients=False, model_factory=None)` `classmethod` ¶

`from_pretrained(repo_name, env, access_token=None, version=None, load_gradients=False, model_factory=None)` `classmethod` ¶

`publish_to_hub(repo_name, folder_path, access_token=None)` ¶