Skip to content

Commit f3a8855

Browse files
MorvanZhouMorvanZhou
committed
add robot arm
1 parent f1632a7 commit f3a8855

2 files changed

Lines changed: 480 additions & 0 deletions

File tree

‎Reinforcement_learning_TUT/experiments/arm/DDPG.py‎

Lines changed: 258 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,258 @@
1+
import tensorflow as tf
2+
import numpy as np
3+
import os
4+
import shutil
5+
from arm_env import ArmEnv
6+
7+
8+
np.random.seed(1)
9+
tf.set_random_seed(1)
10+
11+
MAX_EPISODES = 600
12+
MAX_EP_STEPS = 200
13+
LR_A = 1e-4 # learning rate for actor
14+
LR_C = 1e-4 # learning rate for critic
15+
GAMMA = 0.999 # reward discount
16+
REPLACE_ITER_A = 1100
17+
REPLACE_ITER_C = 1000
18+
MEMORY_CAPACITY = 10000
19+
BATCH_SIZE = 16
20+
VAR_MIN = 0.1
21+
RENDER = True
22+
LOAD = True
23+
MODE = ['easy', 'hard']
24+
n_model = 1
25+
26+
env = ArmEnv(mode=MODE[n_model])
27+
STATE_DIM = env.state_dim
28+
ACTION_DIM = env.action_dim
29+
ACTION_BOUND = env.action_bound
30+
31+
# all placeholder for tf
32+
with tf.name_scope('S'):
33+
S = tf.placeholder(tf.float32, shape=[None, STATE_DIM], name='s')
34+
with tf.name_scope('A'):
35+
A = tf.placeholder(tf.float32, shape=[None, ACTION_DIM], name='a')
36+
with tf.name_scope('R'):
37+
R = tf.placeholder(tf.float32, [None, 1], name='r')
38+
with tf.name_scope('S_'):
39+
S_ = tf.placeholder(tf.float32, shape=[None, STATE_DIM], name='s_')
40+
41+
42+
class Actor(object):
43+
def __init__(self, sess, action_dim, action_bound, learning_rate, t_replace_iter):
44+
self.sess = sess
45+
self.a_dim = action_dim
46+
self.action_bound = action_bound
47+
self.lr = learning_rate
48+
self.t_replace_iter = t_replace_iter
49+
self.t_replace_counter = 0
50+
51+
with tf.variable_scope('Actor'):
52+
# input s, output a
53+
self.a = self._build_net(S, scope='eval_net', trainable=True)
54+
55+
# input s_, output a, get a_ for critic
56+
self.a_ = self._build_net(S_, scope='target_net', trainable=False)
57+
58+
self.e_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='Actor/eval_net')
59+
self.t_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='Actor/target_net')
60+
61+
def _build_net(self, s, scope, trainable):
62+
with tf.variable_scope(scope):
63+
init_w = tf.contrib.layers.xavier_initializer()
64+
init_b = tf.constant_initializer(0.001)
65+
net = tf.layers.dense(s, 200, activation=tf.nn.relu6,
66+
kernel_initializer=init_w, bias_initializer=init_b, name='l1',
67+
trainable=trainable)
68+
net = tf.layers.dense(net, 200, activation=tf.nn.relu6,
69+
kernel_initializer=init_w, bias_initializer=init_b, name='l2',
70+
trainable=trainable)
71+
net = tf.layers.dense(net, 10, activation=tf.nn.relu,
72+
kernel_initializer=init_w, bias_initializer=init_b, name='l3',
73+
trainable=trainable)
74+
with tf.variable_scope('a'):
75+
actions = tf.layers.dense(net, self.a_dim, activation=tf.nn.tanh, kernel_initializer=init_w,
76+
name='a', trainable=trainable)
77+
scaled_a = tf.multiply(actions, self.action_bound, name='scaled_a') # Scale output to -action_bound to action_bound
78+
return scaled_a
79+
80+
def learn(self, s, a): # batch update
81+
self.sess.run(self.train_op, feed_dict={S: s, A: a})
82+
if self.t_replace_counter % self.t_replace_iter == 0:
83+
self.sess.run([tf.assign(t, e) for t, e in zip(self.t_params, self.e_params)])
84+
self.t_replace_counter += 1
85+
86+
def choose_action(self, s):
87+
s = s[np.newaxis, :] # single state
88+
return self.sess.run(self.a, feed_dict={S: s})[0] # single action
89+
90+
def add_grad_to_graph(self, a_grads):
91+
with tf.variable_scope('policy_grads'):
92+
self.policy_grads = tf.gradients(ys=self.a, xs=self.e_params, grad_ys=a_grads)
93+
94+
with tf.variable_scope('A_train'):
95+
opt = tf.train.RMSPropOptimizer(-self.lr / BATCH_SIZE) # (- learning rate) for ascent policy, div to take mean
96+
self.train_op = opt.apply_gradients(zip(self.policy_grads, self.e_params))
97+
98+
99+
class Critic(object):
100+
def __init__(self, sess, state_dim, action_dim, learning_rate, gamma, t_replace_iter, a_):
101+
self.sess = sess
102+
self.s_dim = state_dim
103+
self.a_dim = action_dim
104+
self.lr = learning_rate
105+
self.gamma = gamma
106+
self.t_replace_iter = t_replace_iter
107+
self.t_replace_counter = 0
108+
109+
with tf.variable_scope('Critic'):
110+
# Input (s, a), output q
111+
self.q = self._build_net(S, A, 'eval_net', trainable=True)
112+
113+
# Input (s_, a_), output q_ for q_target
114+
self.q_ = self._build_net(S_, a_, 'target_net', trainable=False) # target_q is based on a_ from Actor's target_net
115+
116+
self.e_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='Critic/eval_net')
117+
self.t_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='Critic/target_net')
118+
119+
with tf.variable_scope('target_q'):
120+
self.target_q = R + self.gamma * self.q_
121+
122+
with tf.variable_scope('TD_error'):
123+
self.loss = tf.reduce_mean(tf.squared_difference(self.target_q, self.q))
124+
125+
with tf.variable_scope('C_train'):
126+
self.train_op = tf.train.RMSPropOptimizer(self.lr).minimize(self.loss)
127+
128+
with tf.variable_scope('a_grad'):
129+
self.a_grads = tf.gradients(self.q, A)[0] # tensor of gradients of each sample (None, a_dim)
130+
131+
def _build_net(self, s, a, scope, trainable):
132+
with tf.variable_scope(scope):
133+
init_w = tf.contrib.layers.xavier_initializer()
134+
init_b = tf.constant_initializer(0.01)
135+
136+
with tf.variable_scope('l1'):
137+
n_l1 = 200
138+
w1_s = tf.get_variable('w1_s', [self.s_dim, n_l1], initializer=init_w, trainable=trainable)
139+
w1_a = tf.get_variable('w1_a', [self.a_dim, n_l1], initializer=init_w, trainable=trainable)
140+
b1 = tf.get_variable('b1', [1, n_l1], initializer=init_b, trainable=trainable)
141+
net = tf.nn.relu6(tf.matmul(s, w1_s) + tf.matmul(a, w1_a) + b1)
142+
net = tf.layers.dense(net, 200, activation=tf.nn.relu6,
143+
kernel_initializer=init_w, bias_initializer=init_b, name='l2',
144+
trainable=trainable)
145+
net = tf.layers.dense(net, 10, activation=tf.nn.relu,
146+
kernel_initializer=init_w, bias_initializer=init_b, name='l3',
147+
trainable=trainable)
148+
with tf.variable_scope('q'):
149+
q = tf.layers.dense(net, 1, kernel_initializer=init_w, bias_initializer=init_b, trainable=trainable) # Q(s,a)
150+
return q
151+
152+
def learn(self, s, a, r, s_):
153+
self.sess.run(self.train_op, feed_dict={S: s, A: a, R: r, S_: s_})
154+
if self.t_replace_counter % self.t_replace_iter == 0:
155+
self.sess.run([tf.assign(t, e) for t, e in zip(self.t_params, self.e_params)])
156+
self.t_replace_counter += 1
157+
158+
159+
class Memory(object):
160+
def __init__(self, capacity, dims):
161+
self.capacity = capacity
162+
self.data = np.zeros((capacity, dims))
163+
self.pointer = 0
164+
165+
def store_transition(self, s, a, r, s_):
166+
transition = np.hstack((s, a, [r], s_))
167+
index = self.pointer % self.capacity # replace the old memory with new memory
168+
self.data[index, :] = transition
169+
self.pointer += 1
170+
171+
def sample(self, n):
172+
assert self.pointer >= self.capacity, 'Memory has not been fulfilled'
173+
indices = np.random.choice(self.capacity, size=n)
174+
return self.data[indices, :]
175+
176+
177+
sess = tf.Session()
178+
179+
# Create actor and critic.
180+
actor = Actor(sess, ACTION_DIM, ACTION_BOUND[1], LR_A, REPLACE_ITER_A)
181+
critic = Critic(sess, STATE_DIM, ACTION_DIM, LR_C, GAMMA, REPLACE_ITER_C, actor.a_)
182+
actor.add_grad_to_graph(critic.a_grads)
183+
184+
M = Memory(MEMORY_CAPACITY, dims=2 * STATE_DIM + ACTION_DIM + 1)
185+
186+
saver = tf.train.Saver()
187+
path = './'+MODE[n_model]
188+
189+
if LOAD:
190+
saver.restore(sess, tf.train.latest_checkpoint(path))
191+
else:
192+
sess.run(tf.global_variables_initializer())
193+
194+
195+
def train():
196+
var = 2. # control exploration
197+
198+
for ep in range(MAX_EPISODES):
199+
s = env.reset()
200+
ep_reward = 0
201+
202+
for t in range(MAX_EP_STEPS):
203+
# while True:
204+
if RENDER:
205+
env.render()
206+
207+
# Added exploration noise
208+
a = actor.choose_action(s)
209+
a = np.clip(np.random.normal(a, var), *ACTION_BOUND) # add randomness to action selection for exploration
210+
s_, r, done = env.step(a)
211+
M.store_transition(s, a, r, s_)
212+
213+
if M.pointer > MEMORY_CAPACITY:
214+
var = max([var*.99995, VAR_MIN]) # decay the action randomness
215+
b_M = M.sample(BATCH_SIZE)
216+
b_s = b_M[:, :STATE_DIM]
217+
b_a = b_M[:, STATE_DIM: STATE_DIM + ACTION_DIM]
218+
b_r = b_M[:, -STATE_DIM - 1: -STATE_DIM]
219+
b_s_ = b_M[:, -STATE_DIM:]
220+
221+
critic.learn(b_s, b_a, b_r, b_s_)
222+
actor.learn(b_s, b_a)
223+
224+
s = s_
225+
ep_reward += r
226+
227+
if t == MAX_EP_STEPS-1 or done:
228+
# if done:
229+
result = '| done' if done else '| ----'
230+
print('Ep:', ep,
231+
result,
232+
'| R: %i' % int(ep_reward),
233+
'| Explore: %.2f' % var,
234+
)
235+
break
236+
237+
if os.path.isdir(path): shutil.rmtree(path)
238+
os.mkdir(path)
239+
ckpt_path = os.path.join('./'+MODE[n_model], 'DDPG.ckpt')
240+
save_path = saver.save(sess, ckpt_path, write_meta_graph=False)
241+
print("\nSave Model %s\n" % save_path)
242+
243+
244+
def eval():
245+
env.set_fps(30)
246+
s = env.reset()
247+
while True:
248+
if RENDER:
249+
env.render()
250+
a = actor.choose_action(s)
251+
s_, r, done = env.step(a)
252+
s = s_
253+
254+
if __name__ == '__main__':
255+
if LOAD:
256+
eval()
257+
else:
258+
train()

0 commit comments

Comments
 (0)