policy-dynamics-value-functions/pdvf_storage.py at master · rraileanu/policy-dynamics-value-functions · GitHub

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
57
58
from collections import namedtuple
import random


TransitionPDVF = namedtuple('Transition',
                        ('state', 'emb_policy', 'emb_env', 'total_return'))

TransitionPolicyDecoder = namedtuple('TransitionPolicyDecoder',
                        ('emb_state', 'recurrent_state', 'mask', 'action'))


class ReplayMemoryPDVF(object):
    '''
    Storage for training the PDVF.
    '''

    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a transition."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)

        self.memory[self.position] = TransitionPDVF(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)


class ReplayMemoryPolicyDecoder(object):
    '''
    Storage for training the olicy decoder.
    '''
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a transition."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)

        self.memory[self.position] = TransitionPolicyDecoder(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)