mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-08-31 12:11:04 +08:00
35 lines
833 B
Python
35 lines
833 B
Python
# -*- coding: utf-8 -*-
|
|
"""Common util functions for all algorithms.
|
|
|
|
- Author: Curt Park
|
|
- Contact: curt.park@medipixel.io
|
|
"""
|
|
|
|
import random
|
|
|
|
import gym
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
def identity(x: torch.Tensor) -> torch.Tensor:
|
|
"""Return input without any change."""
|
|
return x
|
|
|
|
|
|
def soft_update(local: nn.Module, target: nn.Module, tau: float):
|
|
"""Soft-update: target = tau*local + (1-tau)*target."""
|
|
for t_param, l_param in zip(target.parameters(), local.parameters()):
|
|
t_param.data.copy_(tau * l_param.data + (1.0 - tau) * t_param.data)
|
|
|
|
|
|
def set_random_seed(seed: int, env: gym.Env):
|
|
"""Set random seed"""
|
|
env.seed(seed)
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|