wehub-resource-sync
887 lines
30 KiB
Python
887 lines
30 KiB
Python
import logging
|
|
import os
|
|
import warnings
|
|
from typing import TYPE_CHECKING, Dict, List, Optional, Union
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
import tree # pip install dm_tree
|
|
from gymnasium.spaces import Discrete, MultiDiscrete
|
|
from packaging import version
|
|
|
|
from ray.rllib.models.repeated_values import RepeatedValues
|
|
from ray.rllib.utils.annotations import DeveloperAPI, OldAPIStack, PublicAPI
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
from ray.rllib.utils.numpy import SMALL_NUMBER
|
|
from ray.rllib.utils.typing import (
|
|
LocalOptimizer,
|
|
NetworkType,
|
|
SpaceStruct,
|
|
TensorStructType,
|
|
TensorType,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.rllib.core.learner.learner import ParamDict, ParamList
|
|
from ray.rllib.policy.torch_policy import TorchPolicy
|
|
from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
|
|
|
|
logger = logging.getLogger(__name__)
|
|
torch, nn = try_import_torch()
|
|
|
|
# Limit values suitable for use as close to a -inf logit. These are useful
|
|
# since -inf / inf cause NaNs during backprop.
|
|
FLOAT_MIN = -3.4e38
|
|
FLOAT_MAX = 3.4e38
|
|
|
|
if torch:
|
|
TORCH_COMPILE_REQUIRED_VERSION = version.parse("2.0.0")
|
|
else:
|
|
TORCH_COMPILE_REQUIRED_VERSION = ValueError(
|
|
"torch is not installed. TORCH_COMPILE_REQUIRED_VERSION is not defined."
|
|
)
|
|
|
|
|
|
@OldAPIStack
|
|
def apply_grad_clipping(
|
|
policy: "TorchPolicy", optimizer: LocalOptimizer, loss: TensorType
|
|
) -> Dict[str, TensorType]:
|
|
"""Applies gradient clipping to already computed grads inside `optimizer`.
|
|
|
|
Note: This function does NOT perform an analogous operation as
|
|
tf.clip_by_global_norm. It merely clips by norm (per gradient tensor) and
|
|
then computes the global norm across all given tensors (but without clipping
|
|
by that global norm).
|
|
|
|
Args:
|
|
policy: The TorchPolicy, which calculated `loss`.
|
|
optimizer: A local torch optimizer object.
|
|
loss: The torch loss tensor.
|
|
|
|
Returns:
|
|
An info dict containing the "grad_norm" key and the resulting clipped
|
|
gradients.
|
|
"""
|
|
grad_gnorm = 0
|
|
if policy.config["grad_clip"] is not None:
|
|
clip_value = policy.config["grad_clip"]
|
|
else:
|
|
clip_value = np.inf
|
|
|
|
num_none_grads = 0
|
|
for param_group in optimizer.param_groups:
|
|
# Make sure we only pass params with grad != None into torch
|
|
# clip_grad_norm_. Would fail otherwise.
|
|
params = list(filter(lambda p: p.grad is not None, param_group["params"]))
|
|
if params:
|
|
# PyTorch clips gradients inplace and returns the norm before clipping
|
|
# We therefore need to compute grad_gnorm further down (fixes #4965)
|
|
global_norm = nn.utils.clip_grad_norm_(params, clip_value)
|
|
|
|
if isinstance(global_norm, torch.Tensor):
|
|
global_norm = global_norm.cpu().numpy()
|
|
|
|
grad_gnorm += min(global_norm, clip_value)
|
|
else:
|
|
num_none_grads += 1
|
|
|
|
# Note (Kourosh): grads could indeed be zero. This method should still return
|
|
# grad_gnorm in that case.
|
|
if num_none_grads == len(optimizer.param_groups):
|
|
# No grads available
|
|
return {}
|
|
return {"grad_gnorm": grad_gnorm}
|
|
|
|
|
|
@PublicAPI
|
|
def clip_gradients(
|
|
gradients_dict: "ParamDict",
|
|
*,
|
|
grad_clip: Optional[float] = None,
|
|
grad_clip_by: str = "value",
|
|
) -> TensorType:
|
|
"""Performs gradient clipping on a grad-dict based on a clip value and clip mode.
|
|
|
|
Changes the provided gradient dict in place.
|
|
|
|
Args:
|
|
gradients_dict: The gradients dict, mapping str to gradient tensors.
|
|
grad_clip: The value to clip with. The way gradients are clipped is defined
|
|
by the `grad_clip_by` arg (see below).
|
|
grad_clip_by: One of 'value', 'norm', or 'global_norm'.
|
|
|
|
Returns:
|
|
If `grad_clip_by`="global_norm" and `grad_clip` is not None, returns the global
|
|
norm of all tensors, otherwise returns None.
|
|
"""
|
|
# No clipping, return.
|
|
if grad_clip is None:
|
|
return
|
|
|
|
if grad_clip_by not in ["value", "norm", "global_norm"]:
|
|
raise ValueError(
|
|
f"`grad_clip_by` ({grad_clip_by}) must be one of [value|norm|global_norm]!"
|
|
)
|
|
|
|
# Clip by value (each gradient individually).
|
|
if grad_clip_by == "value":
|
|
for k, v in gradients_dict.items():
|
|
gradients_dict[k] = (
|
|
None if v is None else torch.clip(v, -grad_clip, grad_clip)
|
|
)
|
|
|
|
# Clip by L2-norm (per gradient tensor).
|
|
elif grad_clip_by == "norm":
|
|
for k, v in gradients_dict.items():
|
|
if v is not None:
|
|
# Compute the L2-norm of the gradient tensor.
|
|
norm = v.norm(2).nan_to_num(neginf=-10e8, posinf=10e8)
|
|
# Clip all the gradients.
|
|
if norm > grad_clip:
|
|
v.mul_(grad_clip / norm)
|
|
|
|
# Clip by global L2-norm (across all gradient tensors).
|
|
else:
|
|
gradients_list = list(gradients_dict.values())
|
|
total_norm = compute_global_norm(gradients_list)
|
|
if len(gradients_list) == 0:
|
|
return total_norm
|
|
# We do want the coefficient to be in between 0.0 and 1.0, therefore
|
|
# if the global_norm is smaller than the clip value, we use the clip value
|
|
# as normalization constant.
|
|
clip_coeff = grad_clip / torch.clamp(total_norm + 1e-6, min=grad_clip)
|
|
# Note: multiplying by the clamped coefficient is redundant when the coefficient
|
|
# is clamped to 1, but doing so avoids a `if clip_coeff < 1:` conditional which
|
|
# can require a CPU <=> device synchronization when the gradients reside in GPU
|
|
# memory.
|
|
clip_coeff_clamped = torch.clamp(clip_coeff, max=1.0)
|
|
for g in gradients_list:
|
|
if g is not None:
|
|
g.detach().mul_(clip_coeff_clamped.to(g.device))
|
|
return total_norm
|
|
|
|
|
|
@PublicAPI
|
|
def compute_global_norm(gradients_list: "ParamList") -> TensorType:
|
|
"""Computes the global norm for a gradients dict.
|
|
|
|
Args:
|
|
gradients_list: The gradients list containing parameters.
|
|
|
|
Returns:
|
|
Returns the global norm of all tensors in `gradients_list`.
|
|
"""
|
|
# Define the norm type to be L2.
|
|
norm_type = 2.0
|
|
# If we have no grads, return zero.
|
|
if len(gradients_list) == 0:
|
|
return torch.tensor(0.0)
|
|
|
|
# Compute the global norm.
|
|
total_norm = torch.norm(
|
|
torch.stack(
|
|
[
|
|
torch.norm(g.detach(), norm_type)
|
|
# Note, we want to avoid overflow in the norm computation, this does
|
|
# not affect the gradients themselves as we clamp by multiplying and
|
|
# not by overriding tensor values.
|
|
.nan_to_num(neginf=-10e8, posinf=10e8)
|
|
for g in gradients_list
|
|
if g is not None
|
|
]
|
|
),
|
|
norm_type,
|
|
).nan_to_num(neginf=-10e8, posinf=10e8)
|
|
|
|
# Return the global norm.
|
|
return total_norm
|
|
|
|
|
|
@OldAPIStack
|
|
def concat_multi_gpu_td_errors(
|
|
policy: Union["TorchPolicy", "TorchPolicyV2"]
|
|
) -> Dict[str, TensorType]:
|
|
"""Concatenates multi-GPU (per-tower) TD error tensors given TorchPolicy.
|
|
|
|
TD-errors are extracted from the TorchPolicy via its tower_stats property.
|
|
|
|
Args:
|
|
policy: The TorchPolicy to extract the TD-error values from.
|
|
|
|
Returns:
|
|
A dict mapping strings "td_error" and "mean_td_error" to the
|
|
corresponding concatenated and mean-reduced values.
|
|
"""
|
|
td_error = torch.cat(
|
|
[
|
|
t.tower_stats.get("td_error", torch.tensor([0.0])).to(policy.device)
|
|
for t in policy.model_gpu_towers
|
|
],
|
|
dim=0,
|
|
)
|
|
policy.td_error = td_error
|
|
return {
|
|
"td_error": td_error,
|
|
"mean_td_error": torch.mean(td_error),
|
|
}
|
|
|
|
|
|
@PublicAPI
|
|
def convert_to_torch_tensor(
|
|
x,
|
|
device: Optional[str] = None,
|
|
pin_memory: bool = False,
|
|
use_stream: bool = False,
|
|
stream: Optional[Union["torch.cuda.Stream", "torch.cuda.classes.Stream"]] = None,
|
|
):
|
|
"""
|
|
Converts any (possibly nested) structure to torch.Tensors.
|
|
|
|
Args:
|
|
x: The input structure whose leaves will be converted.
|
|
device: The device to create the tensor on (e.g. "cuda:0" or "cpu").
|
|
pin_memory: If True, calls `pin_memory()` on the created tensors.
|
|
use_stream: If True, uses a separate CUDA stream for `Tensor.to()`.
|
|
stream: An optional CUDA stream for the host-to-device copy in `Tensor.to()`.
|
|
|
|
Returns:
|
|
A new structure with the same layout as `x` but with all leaves converted
|
|
to torch.Tensors. Leaves that are None are left unchanged.
|
|
"""
|
|
|
|
# Convert the provided device (if any) to a torch.device; default to CPU.
|
|
device = torch.device(device) if device is not None else torch.device("cpu")
|
|
is_cuda = (device.type == "cuda") and torch.cuda.is_available()
|
|
|
|
# Determine the appropriate stream.
|
|
if is_cuda:
|
|
if use_stream:
|
|
if stream is not None:
|
|
# Ensure the provided stream is of an acceptable type.
|
|
assert isinstance(
|
|
stream, (torch.cuda.Stream, torch.cuda.classes.Stream)
|
|
), f"`stream` must be a torch.cuda.Stream but got {type(stream)}."
|
|
else:
|
|
stream = torch.cuda.Stream()
|
|
else:
|
|
stream = torch.cuda.default_stream(device=device)
|
|
else:
|
|
stream = None
|
|
|
|
def mapping(item):
|
|
# Pass through None values.
|
|
if item is None:
|
|
return item
|
|
|
|
# Special handling for "RepeatedValues" types.
|
|
if isinstance(item, RepeatedValues):
|
|
return RepeatedValues(
|
|
tree.map_structure(mapping, item.values),
|
|
item.lengths,
|
|
item.max_len,
|
|
)
|
|
|
|
# Convert to a tensor if not already one.
|
|
if torch.is_tensor(item):
|
|
tensor = item
|
|
elif isinstance(item, np.ndarray):
|
|
# Leave object or string arrays as is.
|
|
if item.dtype == object or item.dtype.type is np.str_:
|
|
return item
|
|
# If the numpy array is not writable, suppress warnings.
|
|
if not item.flags.writeable:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
tensor = torch.from_numpy(item)
|
|
else:
|
|
tensor = torch.from_numpy(item)
|
|
else:
|
|
tensor = torch.from_numpy(np.asarray(item))
|
|
|
|
# Convert floating-point tensors from float64 to float32 (unless they are float16).
|
|
if tensor.is_floating_point() and tensor.dtype != torch.float16:
|
|
tensor = tensor.float()
|
|
|
|
# Optionally pin memory for faster host-to-GPU copies.
|
|
if pin_memory and is_cuda:
|
|
tensor = tensor.pin_memory()
|
|
|
|
# Move the tensor to the desired device.
|
|
# For CUDA devices, use the provided stream context if available.
|
|
if is_cuda:
|
|
if stream is not None:
|
|
with torch.cuda.stream(stream):
|
|
tensor = tensor.to(device, non_blocking=True)
|
|
else:
|
|
tensor = tensor.to(device, non_blocking=True)
|
|
else:
|
|
# For CPU (or non-CUDA), this is a no-op if already on the target device.
|
|
tensor = tensor.to(device)
|
|
|
|
return tensor
|
|
|
|
return tree.map_structure(mapping, x)
|
|
|
|
|
|
@PublicAPI
|
|
def copy_torch_tensors(x: TensorStructType, device: Optional[str] = None):
|
|
"""Creates a copy of `x` and makes deep copies torch.Tensors in x.
|
|
|
|
Also moves the copied tensors to the specified device (if not None).
|
|
|
|
Note if an object in x is not a torch.Tensor, it will be shallow-copied.
|
|
|
|
Args:
|
|
x : Any (possibly nested) struct possibly containing torch.Tensors.
|
|
device : The device to move the tensors to.
|
|
|
|
Returns:
|
|
Any: A new struct with the same structure as `x`, but with all
|
|
torch.Tensors deep-copied and moved to the specified device.
|
|
|
|
"""
|
|
|
|
def mapping(item):
|
|
if isinstance(item, torch.Tensor):
|
|
return (
|
|
torch.clone(item.detach())
|
|
if device is None
|
|
else item.detach().to(device)
|
|
)
|
|
else:
|
|
return item
|
|
|
|
return tree.map_structure(mapping, x)
|
|
|
|
|
|
@PublicAPI
|
|
def explained_variance(y: TensorType, pred: TensorType) -> TensorType:
|
|
"""Computes the explained variance for a pair of labels and predictions.
|
|
|
|
The formula used is:
|
|
max(-1.0, 1.0 - (std(y - pred)^2 / std(y)^2))
|
|
|
|
Args:
|
|
y: The labels.
|
|
pred: The predictions.
|
|
|
|
Returns:
|
|
The explained variance given a pair of labels and predictions.
|
|
"""
|
|
squeezed_y = y.squeeze()
|
|
y_var = torch.var(squeezed_y, dim=0)
|
|
diff_var = torch.var(squeezed_y - pred.squeeze(), dim=0)
|
|
min_ = torch.tensor([-1.0]).to(pred.device)
|
|
return torch.max(min_, 1 - (diff_var / (y_var + SMALL_NUMBER)))[0]
|
|
|
|
|
|
@PublicAPI
|
|
def flatten_inputs_to_1d_tensor(
|
|
inputs: TensorStructType,
|
|
spaces_struct: Optional[SpaceStruct] = None,
|
|
time_axis: bool = False,
|
|
) -> TensorType:
|
|
"""Flattens arbitrary input structs according to the given spaces struct.
|
|
|
|
Returns a single 1D tensor resulting from the different input
|
|
components' values.
|
|
|
|
Thereby:
|
|
- Boxes (any shape) get flattened to (B, [T]?, -1). Note that image boxes
|
|
are not treated differently from other types of Boxes and get
|
|
flattened as well.
|
|
- Discrete (int) values are one-hot'd, e.g. a batch of [1, 0, 3] (B=3 with
|
|
Discrete(4) space) results in [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1]].
|
|
- MultiDiscrete values are multi-one-hot'd, e.g. a batch of
|
|
[[0, 2], [1, 4]] (B=2 with MultiDiscrete([2, 5]) space) results in
|
|
[[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 0, 1]].
|
|
|
|
Args:
|
|
inputs: The inputs to be flattened.
|
|
spaces_struct: The structure of the spaces that behind the input
|
|
time_axis: Whether all inputs have a time-axis (after the batch axis).
|
|
If True, will keep not only the batch axis (0th), but the time axis
|
|
(1st) as-is and flatten everything from the 2nd axis up.
|
|
|
|
Returns:
|
|
A single 1D tensor resulting from concatenating all
|
|
flattened/one-hot'd input components. Depending on the time_axis flag,
|
|
the shape is (B, n) or (B, T, n).
|
|
|
|
.. testcode::
|
|
|
|
from gymnasium.spaces import Discrete, Box
|
|
from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor
|
|
import torch
|
|
struct = {
|
|
"a": np.array([1, 3]),
|
|
"b": (
|
|
np.array([[1.0, 2.0], [4.0, 5.0]]),
|
|
np.array(
|
|
[[[8.0], [7.0]], [[5.0], [4.0]]]
|
|
),
|
|
),
|
|
"c": {
|
|
"cb": np.array([1.0, 2.0]),
|
|
},
|
|
}
|
|
struct_torch = tree.map_structure(lambda s: torch.from_numpy(s), struct)
|
|
spaces = dict(
|
|
{
|
|
"a": gym.spaces.Discrete(4),
|
|
"b": (gym.spaces.Box(-1.0, 10.0, (2,)), gym.spaces.Box(-1.0, 1.0, (2,
|
|
1))),
|
|
"c": dict(
|
|
{
|
|
"cb": gym.spaces.Box(-1.0, 1.0, ()),
|
|
}
|
|
),
|
|
}
|
|
)
|
|
print(flatten_inputs_to_1d_tensor(struct_torch, spaces_struct=spaces))
|
|
|
|
.. testoutput::
|
|
|
|
tensor([[0., 1., 0., 0., 1., 2., 8., 7., 1.],
|
|
[0., 0., 0., 1., 4., 5., 5., 4., 2.]])
|
|
|
|
"""
|
|
|
|
flat_inputs = tree.flatten(inputs)
|
|
flat_spaces = (
|
|
tree.flatten(spaces_struct)
|
|
if spaces_struct is not None
|
|
else [None] * len(flat_inputs)
|
|
)
|
|
|
|
B = None
|
|
T = None
|
|
out = []
|
|
for input_, space in zip(flat_inputs, flat_spaces):
|
|
# Store batch and (if applicable) time dimension.
|
|
if B is None:
|
|
B = input_.shape[0]
|
|
if time_axis:
|
|
T = input_.shape[1]
|
|
|
|
# One-hot encoding.
|
|
if isinstance(space, Discrete):
|
|
if time_axis:
|
|
input_ = torch.reshape(input_, [B * T])
|
|
out.append(one_hot(input_, space).float())
|
|
# Multi one-hot encoding.
|
|
elif isinstance(space, MultiDiscrete):
|
|
if time_axis:
|
|
input_ = torch.reshape(input_, [B * T, -1])
|
|
out.append(one_hot(input_, space).float())
|
|
# Box: Flatten.
|
|
else:
|
|
if time_axis:
|
|
input_ = torch.reshape(input_, [B * T, -1])
|
|
else:
|
|
input_ = torch.reshape(input_, [B, -1])
|
|
out.append(input_.float())
|
|
|
|
merged = torch.cat(out, dim=-1)
|
|
# Restore the time-dimension, if applicable.
|
|
if time_axis:
|
|
merged = torch.reshape(merged, [B, T, -1])
|
|
|
|
return merged
|
|
|
|
|
|
@PublicAPI
|
|
def global_norm(tensors: List[TensorType]) -> TensorType:
|
|
"""Returns the global L2 norm over a list of tensors.
|
|
|
|
output = sqrt(SUM(t ** 2 for t in tensors)),
|
|
where SUM reduces over all tensors and over all elements in tensors.
|
|
|
|
Args:
|
|
tensors: The list of tensors to calculate the global norm over.
|
|
|
|
Returns:
|
|
The global L2 norm over the given tensor list.
|
|
"""
|
|
# List of single tensors' L2 norms: SQRT(SUM(xi^2)) over all xi in tensor.
|
|
single_l2s = [torch.pow(torch.sum(torch.pow(t, 2.0)), 0.5) for t in tensors]
|
|
# Compute global norm from all single tensors' L2 norms.
|
|
return torch.pow(sum(torch.pow(l2, 2.0) for l2 in single_l2s), 0.5)
|
|
|
|
|
|
@OldAPIStack
|
|
def huber_loss(x: TensorType, delta: float = 1.0) -> TensorType:
|
|
"""Computes the huber loss for a given term and delta parameter.
|
|
|
|
Reference: https://en.wikipedia.org/wiki/Huber_loss
|
|
Note that the factor of 0.5 is implicitly included in the calculation.
|
|
|
|
Formula:
|
|
L = 0.5 * x^2 for small abs x (delta threshold)
|
|
L = delta * (abs(x) - 0.5*delta) for larger abs x (delta threshold)
|
|
|
|
Args:
|
|
x: The input term, e.g. a TD error.
|
|
delta: The delta parmameter in the above formula.
|
|
|
|
Returns:
|
|
The Huber loss resulting from `x` and `delta`.
|
|
"""
|
|
return torch.where(
|
|
torch.abs(x) < delta,
|
|
torch.pow(x, 2.0) * 0.5,
|
|
delta * (torch.abs(x) - 0.5 * delta),
|
|
)
|
|
|
|
|
|
@OldAPIStack
|
|
def l2_loss(x: TensorType) -> TensorType:
|
|
"""Computes half the L2 norm over a tensor's values without the sqrt.
|
|
|
|
output = 0.5 * sum(x ** 2)
|
|
|
|
Args:
|
|
x: The input tensor.
|
|
|
|
Returns:
|
|
0.5 times the L2 norm over the given tensor's values (w/o sqrt).
|
|
"""
|
|
return 0.5 * torch.sum(torch.pow(x, 2.0))
|
|
|
|
|
|
@PublicAPI
|
|
def one_hot(x: TensorType, space: gym.Space) -> TensorType:
|
|
"""Returns a one-hot tensor, given and int tensor and a space.
|
|
|
|
Handles the MultiDiscrete case as well.
|
|
|
|
Args:
|
|
x: The input tensor.
|
|
space: The space to use for generating the one-hot tensor.
|
|
|
|
Returns:
|
|
The resulting one-hot tensor.
|
|
|
|
Raises:
|
|
ValueError: If the given space is not a discrete one.
|
|
|
|
.. testcode::
|
|
|
|
import torch
|
|
import gymnasium as gym
|
|
from ray.rllib.utils.torch_utils import one_hot
|
|
x = torch.IntTensor([0, 3]) # batch-dim=2
|
|
# Discrete space with 4 (one-hot) slots per batch item.
|
|
s = gym.spaces.Discrete(4)
|
|
print(one_hot(x, s))
|
|
x = torch.IntTensor([[0, 1, 2, 3]]) # batch-dim=1
|
|
# MultiDiscrete space with 5 + 4 + 4 + 7 = 20 (one-hot) slots
|
|
# per batch item.
|
|
s = gym.spaces.MultiDiscrete([5, 4, 4, 7])
|
|
print(one_hot(x, s))
|
|
|
|
.. testoutput::
|
|
|
|
tensor([[1, 0, 0, 0],
|
|
[0, 0, 0, 1]])
|
|
tensor([[1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0]])
|
|
"""
|
|
if isinstance(space, Discrete):
|
|
return nn.functional.one_hot(x.long(), space.n)
|
|
elif isinstance(space, MultiDiscrete):
|
|
if isinstance(space.nvec[0], np.ndarray):
|
|
nvec = np.ravel(space.nvec)
|
|
x = x.reshape(x.shape[0], -1)
|
|
else:
|
|
nvec = space.nvec
|
|
return torch.cat(
|
|
[nn.functional.one_hot(x[:, i].long(), n) for i, n in enumerate(nvec)],
|
|
dim=-1,
|
|
)
|
|
else:
|
|
raise ValueError("Unsupported space for `one_hot`: {}".format(space))
|
|
|
|
|
|
@PublicAPI
|
|
def reduce_mean_ignore_inf(x: TensorType, axis: Optional[int] = None) -> TensorType:
|
|
"""Same as torch.mean() but ignores -inf values.
|
|
|
|
Args:
|
|
x: The input tensor to reduce mean over.
|
|
axis: The axis over which to reduce. None for all axes.
|
|
|
|
Returns:
|
|
The mean reduced inputs, ignoring inf values.
|
|
"""
|
|
mask = torch.ne(x, float("-inf"))
|
|
x_zeroed = torch.where(mask, x, torch.zeros_like(x))
|
|
return torch.sum(x_zeroed, axis) / torch.sum(mask.float(), axis)
|
|
|
|
|
|
@PublicAPI
|
|
def sequence_mask(
|
|
lengths: TensorType,
|
|
maxlen: Optional[int] = None,
|
|
dtype=None,
|
|
time_major: bool = False,
|
|
) -> TensorType:
|
|
"""Offers same behavior as tf.sequence_mask for torch.
|
|
|
|
Thanks to Dimitris Papatheodorou
|
|
(https://discuss.pytorch.org/t/pytorch-equivalent-for-tf-sequence-mask/
|
|
39036).
|
|
|
|
Args:
|
|
lengths: The tensor of individual lengths to mask by.
|
|
maxlen: The maximum length to use for the time axis. If None, use
|
|
the max of `lengths`.
|
|
dtype: The torch dtype to use for the resulting mask.
|
|
time_major: Whether to return the mask as [B, T] (False; default) or
|
|
as [T, B] (True).
|
|
|
|
Returns:
|
|
The sequence mask resulting from the given input and parameters.
|
|
"""
|
|
# If maxlen not given, use the longest lengths in the `lengths` tensor.
|
|
if maxlen is None:
|
|
maxlen = lengths.max()
|
|
|
|
mask = torch.ones(tuple(lengths.shape) + (maxlen,))
|
|
|
|
mask = ~(mask.to(lengths.device).cumsum(dim=1).t() > lengths)
|
|
# Time major transformation.
|
|
if not time_major:
|
|
mask = mask.t()
|
|
|
|
# By default, set the mask to be boolean.
|
|
mask.type(dtype or torch.bool)
|
|
|
|
return mask
|
|
|
|
|
|
@PublicAPI
|
|
def update_target_network(
|
|
main_net: NetworkType,
|
|
target_net: NetworkType,
|
|
tau: float,
|
|
) -> None:
|
|
"""Updates a torch.nn.Module target network using Polyak averaging.
|
|
|
|
.. code-block:: text
|
|
|
|
new_target_net_weight = (
|
|
tau * main_net_weight + (1.0 - tau) * current_target_net_weight
|
|
)
|
|
|
|
Args:
|
|
main_net: The nn.Module to update from.
|
|
target_net: The target network to update.
|
|
tau: The tau value to use in the Polyak averaging formula.
|
|
"""
|
|
# Get the current parameters from the Q network.
|
|
state_dict = main_net.state_dict()
|
|
# Use here Polyak averaging.
|
|
new_state_dict = {
|
|
k: tau * state_dict[k] + (1 - tau) * v
|
|
for k, v in target_net.state_dict().items()
|
|
}
|
|
# Apply the new parameters to the target Q network.
|
|
target_net.load_state_dict(new_state_dict)
|
|
|
|
|
|
@DeveloperAPI
|
|
def warn_if_infinite_kl_divergence(
|
|
policy: "TorchPolicy",
|
|
kl_divergence: TensorType,
|
|
) -> None:
|
|
if policy.loss_initialized() and kl_divergence.isinf():
|
|
logger.warning(
|
|
"KL divergence is non-finite, this will likely destabilize your model and"
|
|
" the training process. Action(s) in a specific state have near-zero"
|
|
" probability. This can happen naturally in deterministic environments"
|
|
" where the optimal policy has zero mass for a specific action. To fix this"
|
|
" issue, consider setting the coefficient for the KL loss term to zero or"
|
|
" increasing policy entropy."
|
|
)
|
|
|
|
|
|
@PublicAPI
|
|
def set_torch_seed(seed: Optional[int] = None) -> None:
|
|
"""Sets the torch random seed to the given value.
|
|
|
|
Args:
|
|
seed: The seed to use or None for no seeding.
|
|
"""
|
|
if seed is not None and torch:
|
|
torch.manual_seed(seed)
|
|
# See https://github.com/pytorch/pytorch/issues/47672.
|
|
cuda_version = torch.version.cuda
|
|
if cuda_version is not None and float(torch.version.cuda) >= 10.2:
|
|
# See https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility.
|
|
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
torch.cuda.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed) # if using multi-GPU
|
|
else:
|
|
if version.Version(torch.__version__) >= version.Version("1.8.0"):
|
|
# Not all Operations support this.
|
|
torch.use_deterministic_algorithms(True)
|
|
else:
|
|
torch.set_deterministic(True)
|
|
# This is only for Convolution no problem.
|
|
torch.backends.cudnn.deterministic = True
|
|
# For benchmark=True, CuDNN may choose different algorithms depending on runtime
|
|
# conditions or slight differences in input sizes, even if the seed is fixed,
|
|
# which breaks determinism.
|
|
torch.backends.cudnn.benchmark = False
|
|
|
|
|
|
@PublicAPI
|
|
def softmax_cross_entropy_with_logits(
|
|
logits: TensorType,
|
|
labels: TensorType,
|
|
) -> TensorType:
|
|
"""Same behavior as tf.nn.softmax_cross_entropy_with_logits.
|
|
|
|
Args:
|
|
x: The input predictions.
|
|
labels: The labels corresponding to `x`.
|
|
|
|
Returns:
|
|
The resulting softmax cross-entropy given predictions and labels.
|
|
"""
|
|
return torch.sum(-labels * nn.functional.log_softmax(logits, -1), -1)
|
|
|
|
|
|
@PublicAPI
|
|
def symlog(x: "torch.Tensor") -> "torch.Tensor":
|
|
"""The symlog function as described in [1]:
|
|
|
|
[1] Mastering Diverse Domains through World Models - 2023
|
|
D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
|
|
https://arxiv.org/pdf/2301.04104v1.pdf
|
|
"""
|
|
return torch.sign(x) * torch.log(torch.abs(x) + 1)
|
|
|
|
|
|
@PublicAPI
|
|
def inverse_symlog(y: "torch.Tensor") -> "torch.Tensor":
|
|
"""Inverse of the `symlog` function as desribed in [1]:
|
|
|
|
[1] Mastering Diverse Domains through World Models - 2023
|
|
D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
|
|
https://arxiv.org/pdf/2301.04104v1.pdf
|
|
"""
|
|
# To get to symlog inverse, we solve the symlog equation for x:
|
|
# y = sign(x) * log(|x| + 1)
|
|
# <=> y / sign(x) = log(|x| + 1)
|
|
# <=> y = log( x + 1) V x >= 0
|
|
# -y = log(-x + 1) V x < 0
|
|
# <=> exp(y) = x + 1 V x >= 0
|
|
# exp(-y) = -x + 1 V x < 0
|
|
# <=> exp(y) - 1 = x V x >= 0
|
|
# exp(-y) - 1 = -x V x < 0
|
|
# <=> exp(y) - 1 = x V x >= 0 (if x >= 0, then y must also be >= 0)
|
|
# -exp(-y) - 1 = x V x < 0 (if x < 0, then y must also be < 0)
|
|
# <=> sign(y) * (exp(|y|) - 1) = x
|
|
return torch.sign(y) * (torch.exp(torch.abs(y)) - 1)
|
|
|
|
|
|
@PublicAPI
|
|
def two_hot(
|
|
value: "torch.Tensor",
|
|
num_buckets: int = 255,
|
|
lower_bound: float = -20.0,
|
|
upper_bound: float = 20.0,
|
|
device: Optional[str] = None,
|
|
):
|
|
"""Returns a two-hot vector of dim=num_buckets with two entries that are non-zero.
|
|
|
|
See [1] for more details:
|
|
[1] Mastering Diverse Domains through World Models - 2023
|
|
D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
|
|
https://arxiv.org/pdf/2301.04104v1.pdf
|
|
|
|
Entries in the vector represent equally sized buckets within some fixed range
|
|
(`lower_bound` to `upper_bound`).
|
|
Those entries not 0.0 at positions k and k+1 encode the actual `value` and sum
|
|
up to 1.0. They are the weights multiplied by the buckets values at k and k+1 for
|
|
retrieving `value`.
|
|
|
|
Example:
|
|
num_buckets=11
|
|
lower_bound=-5
|
|
upper_bound=5
|
|
value=2.5
|
|
-> [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.0, 0.0]
|
|
-> [-5 -4 -3 -2 -1 0 1 2 3 4 5] (0.5*2 + 0.5*3=2.5)
|
|
|
|
Example:
|
|
num_buckets=5
|
|
lower_bound=-1
|
|
upper_bound=1
|
|
value=0.1
|
|
-> [0.0, 0.0, 0.8, 0.2, 0.0]
|
|
-> [-1 -0.5 0 0.5 1] (0.2*0.5 + 0.8*0=0.1)
|
|
|
|
Args:
|
|
value: The input tensor of shape (B,) to be two-hot encoded.
|
|
num_buckets: The number of buckets to two-hot encode into.
|
|
lower_bound: The lower bound value used for the encoding. If input values are
|
|
lower than this boundary, they will be encoded as `lower_bound`.
|
|
upper_bound: The upper bound value used for the encoding. If input values are
|
|
higher than this boundary, they will be encoded as `upper_bound`.
|
|
|
|
Returns:
|
|
The two-hot encoded tensor of shape (B, num_buckets).
|
|
"""
|
|
# First make sure, values are clipped.
|
|
value = torch.clamp(value, lower_bound, upper_bound)
|
|
# Tensor of batch indices: [0, B=batch size).
|
|
batch_indices = torch.arange(0, value.shape[0], device=device).float()
|
|
# Calculate the step deltas (how much space between each bucket's central value?).
|
|
bucket_delta = (upper_bound - lower_bound) / (num_buckets - 1)
|
|
# Compute the float indices (might be non-int numbers: sitting between two buckets).
|
|
idx = (-lower_bound + value) / bucket_delta
|
|
# k
|
|
k = torch.floor(idx)
|
|
# k+1
|
|
kp1 = torch.ceil(idx)
|
|
# In case k == kp1 (idx is exactly on the bucket boundary), move kp1 up by 1.0.
|
|
# Otherwise, this would result in a NaN in the returned two-hot tensor.
|
|
kp1 = torch.where(k.eq(kp1), kp1 + 1.0, kp1)
|
|
# Iff `kp1` is one beyond our last index (because incoming value is larger than
|
|
# `upper_bound`), move it to one before k (kp1's weight is going to be 0.0 anyways,
|
|
# so it doesn't matter where it points to; we are just avoiding an index error
|
|
# with this).
|
|
kp1 = torch.where(kp1.eq(num_buckets), kp1 - 2.0, kp1)
|
|
# The actual values found at k and k+1 inside the set of buckets.
|
|
values_k = lower_bound + k * bucket_delta
|
|
values_kp1 = lower_bound + kp1 * bucket_delta
|
|
# Compute the two-hot weights (adding up to 1.0) to use at index k and k+1.
|
|
weights_k = (value - values_kp1) / (values_k - values_kp1)
|
|
weights_kp1 = 1.0 - weights_k
|
|
# Compile a tensor of full paths (indices from batch index to feature index) to
|
|
# use for the scatter_nd op.
|
|
indices_k = torch.stack([batch_indices, k], dim=-1)
|
|
indices_kp1 = torch.stack([batch_indices, kp1], dim=-1)
|
|
indices = torch.cat([indices_k, indices_kp1], dim=0).long()
|
|
# The actual values (weights adding up to 1.0) to place at the computed indices.
|
|
updates = torch.cat([weights_k, weights_kp1], dim=0)
|
|
# Call the actual scatter update op, returning a zero-filled tensor, only changed
|
|
# at the given indices.
|
|
output = torch.zeros(value.shape[0], num_buckets, device=device)
|
|
# Set our two-hot values at computed indices.
|
|
output[indices[:, 0], indices[:, 1]] = updates
|
|
return output
|
|
|
|
|
|
def _dynamo_is_available():
|
|
# This only works if torch._dynamo is available
|
|
try:
|
|
# TODO(Artur): Remove this once torch._dynamo is available on CI
|
|
import torch._dynamo as dynamo # noqa: F401
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|