Pool/pool.py at main · NoteDance/Pool · 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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import numpy as np
import multiprocessing as mp
import math


class Pool:
    def __init__(self, env, processes, pool_size, num_steps=None, window_size=None, clearing_freq=None, window_size_=None, random=False):
        self.env = env
        self.processes = processes
        self.pool_size = pool_size
        self.num_steps = num_steps
        self.window_size = window_size
        self.clearing_freq = clearing_freq
        self.window_size_ = window_size_
        self.random = random
        manager=mp.Manager()
        self.state_pool_list=manager.list()
        self.action_pool_list=manager.list()
        self.next_state_pool_list=manager.list()
        self.reward_pool_list=manager.list()
        self.done_pool_list=manager.list()
        if self.clearing_freq!=None:
            self.store_counter=manager.list()
        for _ in range(processes):
            self.state_pool_list.append(None)
            self.action_pool_list.append(None)
            self.next_state_pool_list.append(None)
            self.reward_pool_list.append(None)
            self.done_pool_list.append(None)
            self.store_counter.append(0)
        if random:
            self.inverse_len=manager.list([0 for _ in range(processes)])
            self.lock_list=[mp.Lock() for _ in range(self.processes)]
        else:
            self.lock_list=None

    def pool(self,s,a,next_s,r,done,index=None):
        if self.state_pool_list[index] is None:
            self.state_pool_list[index]=s
            self.action_pool_list[index]=np.expand_dims(a,axis=0)
            self.next_state_pool_list[index]=np.expand_dims(next_s,axis=0)
            self.reward_pool_list[index]=np.expand_dims(r,axis=0)
            self.done_pool_list[index]=np.expand_dims(done,axis=0)
        else:
            self.state_pool_list[index]=np.concatenate((self.state_pool_list[index],s),0)
            self.action_pool_list[index]=np.concatenate((self.action_pool_list[index],np.expand_dims(a,axis=0)),0)
            self.next_state_pool_list[index]=np.concatenate((self.next_state_pool_list[index],np.expand_dims(next_s,axis=0)),0)
            self.reward_pool_list[index]=np.concatenate((self.reward_pool_list[index],np.expand_dims(r,axis=0)),0)
            self.done_pool_list[index]=np.concatenate((self.done_pool[7],np.expand_dims(done,axis=0)),0)
        if self.clearing_freq!=None:
            self.store_counter[index]+=1
            if self.store_counter[index]%self.clearing_freq==0:
                self.state_pool_list[index]=self.state_pool_list[index][self.window_size_:]
                self.action_pool_list[index]=self.action_pool_list[index][self.window_size_:]
                self.next_state_pool_list[index]=self.next_state_pool_list[index][self.window_size_:]
                self.reward_pool_list[index]=self.reward_pool_list[index][self.window_size_:]
                self.done_pool_list[index]=self.done_pool_list[index][self.window_size_:]
        if len(self.state_pool_list[index])>math.ceil(self.pool_size/self.processes):
            if self.window_size!=None:
                self.state_pool_list[index]=self.state_pool_list[index][self.window_size:]
                self.action_pool_list[index]=self.action_pool_list[index][self.window_size:]
                self.next_state_pool_list[index]=self.next_state_pool_list[index][self.window_size:]
                self.reward_pool_list[index]=self.reward_pool_list[index][self.window_size:]
                self.done_pool_list[index]=self.done_pool_list[index][self.window_size:]
            else:
                self.state_pool_list[index]=self.state_pool_list[index][1:]
                self.action_pool_list[index]=self.action_pool_list[index][1:]
                self.next_state_pool_list[index]=self.next_state_pool_list[index][1:]
                self.reward_pool_list[index]=self.reward_pool_list[index][1:]
                self.done_pool_list[index]=self.done_pool_list[index][1:]

    def store_in_parallel(self,env,p,lock_list):
        s,a=env.reset()
        s=np.array(s)
        reward=0
        counter=0
        while True:
            if self.random:
                if self.state_pool_list[p] is None:
                    index=p
                    self.inverse_len[index]=1
                else:
                    total_inverse=np.sum(self.inverse_len)
                    prob=self.inverse_len/total_inverse
                    index=np.random.choice(self.processes,p=prob.numpy(),replace=False)
                    self.inverse_len[index]=1/(len(self.state_pool_list[index])+1)
            else:
                index=p
            s=np.expand_dims(s,axis=0)
            a,next_s,r,done=env.step(a)
            next_s=np.array(next_s)
            r=np.array(r)
            done=np.array(done)
            if self.random:
                lock_list[index].acquire()
                if self.num_steps!=None:
                    if counter==0:
                        next_s_=next_s
                        done_=done
                    counter+=1
                    reward=r+reward
                    if counter%self.num_steps==0 or done:
                        self.pool(s,a,next_s,reward,done,index)
                        reward=0
                else:
                    self.pool(s,a,next_s,r,done,index)
                lock_list[index].release()
            else:
                self.pool(s,a,next_s,r,done,index)
            if (self.num_steps==None and done) or (self.num_steps!=None and done_):
                return
            s=next_s
            if (self.num_steps!=None and counter%self.num_steps==0) or (self.num_steps!=None and done):
                s=next_s_

    def store(self):
        process_list=[]
        for p in range(self.processes):
            process=mp.Process(target=self.store_in_parallel,args=(self.env[p],p,self.lock_list))
            process.start()
            process_list.append(process)
        for process in process_list:
            process.join()

    def get_pool(self):
        state_pool=np.concatenate(self.state_pool_list)
        action_pool=np.concatenate(self.action_pool_list)
        next_state_pool=np.concatenate(self.next_state_pool_list)
        reward_pool=np.concatenate(self.reward_pool_list)
        done_pool=np.concatenate(self.done_pool_list)
        return state_pool, action_pool, next_state_pool, reward_pool, done_pool