You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
52 lines
1.5 KiB
52 lines
1.5 KiB
"""Basic math primitives for Isaac Lab: normalization, random floats, copysign, seeding."""
|
|
|
|
import torch
|
|
import numpy as np
|
|
import random
|
|
import os
|
|
|
|
|
|
@torch.jit.script
|
|
def normalize(x, eps: float = 1e-9):
|
|
return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1)
|
|
|
|
|
|
@torch.jit.script
|
|
def torch_rand_float(lower, upper, shape, device):
|
|
# type: (float, float, Tuple[int, int], str) -> Tensor
|
|
return (upper - lower) * torch.rand(*shape, device=device) + lower
|
|
|
|
|
|
@torch.jit.script
|
|
def copysign(a, b):
|
|
# type: (float, Tensor) -> Tensor
|
|
a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
|
|
return torch.abs(a) * torch.sign(b)
|
|
|
|
|
|
def set_seed(seed, torch_deterministic=False):
|
|
"""set seed across modules"""
|
|
if seed == -1 and torch_deterministic:
|
|
seed = 42
|
|
elif seed == -1:
|
|
seed = np.random.randint(0, 10000)
|
|
print("Setting seed: {}".format(seed))
|
|
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
os.environ["PYTHONHASHSEED"] = str(seed)
|
|
torch.cuda.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
if torch_deterministic:
|
|
# refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
|
|
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
torch.backends.cudnn.benchmark = False
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.use_deterministic_algorithms(True)
|
|
else:
|
|
torch.backends.cudnn.benchmark = True
|
|
torch.backends.cudnn.deterministic = False
|
|
|
|
return seed
|