mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-08-29 11:20:40 +08:00
73 lines
1.9 KiB
Python
Executable File
73 lines
1.9 KiB
Python
Executable File
# -*- coding: utf-8 -*-
|
|
"""Common util functions for all algorithms.
|
|
|
|
- Author: Curt Park
|
|
- Contact: curt.park@medipixel.io
|
|
"""
|
|
|
|
import random
|
|
from collections import deque
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
def identity(x):
|
|
"""Return input without any change."""
|
|
return x
|
|
|
|
|
|
def soft_update(local, target, tau):
|
|
"""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, env):
|
|
"""Set random seed"""
|
|
env.seed(seed)
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
|
|
|
|
def get_n_step_info_from_demo(demo, n_step, gamma):
|
|
"""Return 1 step and n step demos."""
|
|
assert demo
|
|
assert n_step > 1
|
|
|
|
demos_1_step = list()
|
|
demos_n_step = list()
|
|
n_step_buffer = deque(maxlen=n_step)
|
|
|
|
for transition in demo:
|
|
n_step_buffer.append(transition)
|
|
|
|
if len(n_step_buffer) == n_step:
|
|
# add a single step transition
|
|
demos_1_step.append(n_step_buffer[0])
|
|
|
|
# add a multi step transition
|
|
curr_state, action = n_step_buffer[0][:2]
|
|
reward, next_state, done = get_n_step_info(n_step_buffer, gamma)
|
|
transition = (curr_state, action, reward, next_state, done)
|
|
demos_n_step.append(transition)
|
|
|
|
return demos_1_step, demos_n_step
|
|
|
|
|
|
def get_n_step_info(n_step_buffer, gamma):
|
|
"""Return n step reward, next state, and done."""
|
|
# info of the last transition
|
|
reward, next_state, done = n_step_buffer[-1][-3:]
|
|
|
|
for transition in reversed(list(n_step_buffer)[:-1]):
|
|
r, n_s, d = transition[-3:]
|
|
|
|
reward = r + gamma * reward * (1 - d)
|
|
next_state, done = (n_s, d) if d else (next_state, done)
|
|
|
|
return reward, next_state, done
|