166 lines
6.9 KiB
Python
166 lines
6.9 KiB
Python
from typing import Callable, Optional, Sequence
|
||
|
||
import torch as th
|
||
from torch import nn
|
||
|
||
|
||
def quantile_huber_loss(
|
||
current_quantiles: th.Tensor,
|
||
target_quantiles: th.Tensor,
|
||
cum_prob: Optional[th.Tensor] = None,
|
||
sum_over_quantiles: bool = True,
|
||
) -> th.Tensor:
|
||
"""
|
||
The quantile-regression loss, as described in the QR-DQN and TQC papers.
|
||
Partially taken from https://github.com/bayesgroup/tqc_pytorch.
|
||
|
||
:param current_quantiles: current estimate of quantiles, must be either
|
||
(batch_size, n_quantiles) or (batch_size, n_critics, n_quantiles)
|
||
:param target_quantiles: target of quantiles, must be either (batch_size, n_target_quantiles),
|
||
(batch_size, 1, n_target_quantiles), or (batch_size, n_critics, n_target_quantiles)
|
||
:param cum_prob: cumulative probabilities to calculate quantiles (also called midpoints in QR-DQN paper),
|
||
must be either (batch_size, n_quantiles), (batch_size, 1, n_quantiles), or (batch_size, n_critics, n_quantiles).
|
||
(if None, calculating unit quantiles)
|
||
:param sum_over_quantiles: if summing over the quantile dimension or not
|
||
:return: the loss
|
||
"""
|
||
if current_quantiles.ndim != target_quantiles.ndim:
|
||
raise ValueError(
|
||
f"Error: The dimension of curremt_quantile ({current_quantiles.ndim}) needs to match "
|
||
f"the dimension of target_quantiles ({target_quantiles.ndim})."
|
||
)
|
||
if current_quantiles.shape[0] != target_quantiles.shape[0]:
|
||
raise ValueError(
|
||
f"Error: The batch size of curremt_quantile ({current_quantiles.shape[0]}) needs to match "
|
||
f"the batch size of target_quantiles ({target_quantiles.shape[0]})."
|
||
)
|
||
if current_quantiles.ndim not in (2, 3):
|
||
raise ValueError(f"Error: The dimension of current_quantiles ({current_quantiles.ndim}) needs to be either 2 or 3.")
|
||
|
||
if cum_prob is None:
|
||
n_quantiles = current_quantiles.shape[-1]
|
||
# Cumulative probabilities to calculate quantiles.
|
||
cum_prob = (th.arange(n_quantiles, device=current_quantiles.device, dtype=th.float) + 0.5) / n_quantiles
|
||
if current_quantiles.ndim == 2:
|
||
# For QR-DQN, current_quantiles have a shape (batch_size, n_quantiles), and make cum_prob
|
||
# broadcastable to (batch_size, n_quantiles, n_target_quantiles)
|
||
cum_prob = cum_prob.view(1, -1, 1)
|
||
elif current_quantiles.ndim == 3:
|
||
# For TQC, current_quantiles have a shape (batch_size, n_critics, n_quantiles), and make cum_prob
|
||
# broadcastable to (batch_size, n_critics, n_quantiles, n_target_quantiles)
|
||
cum_prob = cum_prob.view(1, 1, -1, 1)
|
||
|
||
# QR-DQN
|
||
# target_quantiles: (batch_size, n_target_quantiles) -> (batch_size, 1, n_target_quantiles)
|
||
# current_quantiles: (batch_size, n_quantiles) -> (batch_size, n_quantiles, 1)
|
||
# pairwise_delta: (batch_size, n_target_quantiles, n_quantiles)
|
||
# TQC
|
||
# target_quantiles: (batch_size, 1, n_target_quantiles) -> (batch_size, 1, 1, n_target_quantiles)
|
||
# current_quantiles: (batch_size, n_critics, n_quantiles) -> (batch_size, n_critics, n_quantiles, 1)
|
||
# pairwise_delta: (batch_size, n_critics, n_quantiles, n_target_quantiles)
|
||
# Note: in both cases, the loss has the same shape as pairwise_delta
|
||
pairwise_delta = target_quantiles.unsqueeze(-2) - current_quantiles.unsqueeze(-1)
|
||
abs_pairwise_delta = th.abs(pairwise_delta)
|
||
huber_loss = th.where(abs_pairwise_delta > 1, abs_pairwise_delta - 0.5, pairwise_delta**2 * 0.5)
|
||
loss = th.abs(cum_prob - (pairwise_delta.detach() < 0).float()) * huber_loss
|
||
if sum_over_quantiles:
|
||
loss = loss.sum(dim=-2).mean()
|
||
else:
|
||
loss = loss.mean()
|
||
return loss
|
||
|
||
|
||
def conjugate_gradient_solver(
|
||
matrix_vector_dot_fn: Callable[[th.Tensor], th.Tensor],
|
||
b,
|
||
max_iter=10,
|
||
residual_tol=1e-10,
|
||
) -> th.Tensor:
|
||
"""
|
||
Finds an approximate solution to a set of linear equations Ax = b
|
||
|
||
Sources:
|
||
- https://github.com/ajlangley/trpo-pytorch/blob/master/conjugate_gradient.py
|
||
- https://github.com/joschu/modular_rl/blob/master/modular_rl/trpo.py#L122
|
||
|
||
Reference:
|
||
- https://epubs.siam.org/doi/abs/10.1137/1.9781611971446.ch6
|
||
|
||
:param matrix_vector_dot_fn:
|
||
a function that right multiplies a matrix A by a vector v
|
||
:param b:
|
||
the right hand term in the set of linear equations Ax = b
|
||
:param max_iter:
|
||
the maximum number of iterations (default is 10)
|
||
:param residual_tol:
|
||
residual tolerance for early stopping of the solving (default is 1e-10)
|
||
:return x:
|
||
the approximate solution to the system of equations defined by `matrix_vector_dot_fn`
|
||
and b
|
||
"""
|
||
|
||
# The vector is not initialized at 0 because of the instability issues when the gradient becomes small.
|
||
# A small random gaussian noise is used for the initialization.
|
||
x = 1e-4 * th.randn_like(b)
|
||
residual = b - matrix_vector_dot_fn(x)
|
||
# Equivalent to th.linalg.norm(residual) ** 2 (L2 norm squared)
|
||
residual_squared_norm = th.matmul(residual, residual)
|
||
|
||
if residual_squared_norm < residual_tol:
|
||
# If the gradient becomes extremely small
|
||
# The denominator in alpha will become zero
|
||
# Leading to a division by zero
|
||
return x
|
||
|
||
p = residual.clone()
|
||
|
||
for i in range(max_iter):
|
||
# A @ p (matrix vector multiplication)
|
||
A_dot_p = matrix_vector_dot_fn(p)
|
||
|
||
alpha = residual_squared_norm / p.dot(A_dot_p)
|
||
x += alpha * p
|
||
|
||
if i == max_iter - 1:
|
||
return x
|
||
|
||
residual -= alpha * A_dot_p
|
||
new_residual_squared_norm = th.matmul(residual, residual)
|
||
|
||
if new_residual_squared_norm < residual_tol:
|
||
return x
|
||
|
||
beta = new_residual_squared_norm / residual_squared_norm
|
||
residual_squared_norm = new_residual_squared_norm
|
||
p = residual + beta * p
|
||
# Note: this return statement is only used when max_iter=0
|
||
return x
|
||
|
||
|
||
def flat_grad(
|
||
output,
|
||
parameters: Sequence[nn.parameter.Parameter],
|
||
create_graph: bool = False,
|
||
retain_graph: bool = False,
|
||
) -> th.Tensor:
|
||
"""
|
||
Returns the gradients of the passed sequence of parameters into a flat gradient.
|
||
Order of parameters is preserved.
|
||
|
||
:param output: functional output to compute the gradient for
|
||
:param parameters: sequence of ``Parameter``
|
||
:param retain_graph: – If ``False``, the graph used to compute the grad will be freed.
|
||
Defaults to the value of ``create_graph``.
|
||
:param create_graph: – If ``True``, graph of the derivative will be constructed,
|
||
allowing to compute higher order derivative products. Default: ``False``.
|
||
:return: Tensor containing the flattened gradients
|
||
"""
|
||
grads = th.autograd.grad(
|
||
output,
|
||
parameters,
|
||
create_graph=create_graph,
|
||
retain_graph=retain_graph,
|
||
allow_unused=True,
|
||
)
|
||
return th.cat([th.ravel(grad) for grad in grads if grad is not None])
|