forked from ryanxhr/CPQ
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
56 lines (44 loc) · 2.01 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import numpy as np
import torch
class ReplayBuffer(object):
def __init__(self, state_dim, action_dim, max_size=int(1e6)):
self.max_size = max_size
self.ptr = 0
self.size = 0
self.state = np.zeros((max_size, state_dim))
self.action = np.zeros((max_size, action_dim))
self.next_state = np.zeros((max_size, state_dim))
self.reward = np.zeros((max_size, 1))
self.not_done = np.zeros((max_size, 1))
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def add(self, state, action, next_state, reward, done):
self.state[self.ptr] = state
self.action[self.ptr] = action
self.next_state[self.ptr] = next_state
self.reward[self.ptr] = reward
self.not_done[self.ptr] = 1. - done
self.ptr = (self.ptr + 1) % self.max_size
self.size = min(self.size + 1, self.max_size)
def sample(self, batch_size):
ind = np.random.randint(0, self.size, size=batch_size)
return (
torch.FloatTensor(self.state[ind]).to(self.device),
torch.FloatTensor(self.action[ind]).to(self.device),
torch.FloatTensor(self.next_state[ind]).to(self.device),
torch.FloatTensor(self.reward[ind]).to(self.device),
torch.FloatTensor(self.not_done[ind]).to(self.device)
)
def convert_D4RL(self, dataset):
self.state = dataset['observations']
self.action = dataset['actions']
self.next_state = dataset['next_observations']
self.reward = dataset['rewards'].reshape(-1, 1)
self.not_done = 1. - dataset['terminals'].reshape(-1, 1)
self.size = self.state.shape[0]
def normalize_states(self, eps=1e-3, mean=None, std=None):
if mean is None and std is None:
mean = self.state.mean(0, keepdims=True)
std = self.state.std(0, keepdims=True) + eps
self.state = (self.state - mean) / std
self.next_state = (self.next_state - mean) / std
return mean, std