-
Notifications
You must be signed in to change notification settings - Fork 817
/
train.py
156 lines (135 loc) · 5.09 KB
/
train.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
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import gym
import argparse
import numpy as np
from parl.utils import logger, summary, ReplayMemory
from parl.env import ActionMappingWrapper, CompatWrapper
from mujoco_model import MujocoModel
from mujoco_agent import MujocoAgent
from parl.algorithms import SAC
WARMUP_STEPS = 1e4
EVAL_EPISODES = 5
MEMORY_SIZE = int(1e6)
BATCH_SIZE = 256
GAMMA = 0.99
TAU = 0.005
ACTOR_LR = 3e-4
CRITIC_LR = 3e-4
# Run episode for training
def run_train_episode(agent, env, rpm):
action_dim = env.action_space.shape[0]
obs = env.reset()
done = False
episode_reward = 0
episode_steps = 0
while not done:
episode_steps += 1
# Select action randomly or according to policy
if rpm.size() < WARMUP_STEPS:
action = np.random.uniform(-1, 1, size=action_dim)
else:
action = agent.sample(obs)
# Perform action
next_obs, reward, done, _ = env.step(action)
terminal = float(done) if episode_steps < env._max_episode_steps else 0
# Store data in replay memory
rpm.append(obs, action, reward, next_obs, terminal)
obs = next_obs
episode_reward += reward
# Train agent after collecting sufficient data
if rpm.size() >= WARMUP_STEPS:
batch_obs, batch_action, batch_reward, batch_next_obs, batch_terminal = rpm.sample_batch(
BATCH_SIZE)
agent.learn(batch_obs, batch_action, batch_reward, batch_next_obs,
batch_terminal)
return episode_reward, episode_steps
# Runs policy for 5 episodes by default and returns average reward
# A fixed seed is used for the eval environment
def run_evaluate_episodes(agent, env, eval_episodes):
avg_reward = 0.
for _ in range(eval_episodes):
obs = env.reset()
done = False
while not done:
action = agent.predict(obs)
obs, reward, done, _ = env.step(action)
avg_reward += reward
avg_reward /= eval_episodes
return avg_reward
def main():
logger.info("------------------- SAC ---------------------")
logger.info('Env: {}, Seed: {}'.format(args.env, args.seed))
logger.info("---------------------------------------------")
logger.set_dir('./{}_{}'.format(args.env, args.seed))
env = gym.make(args.env)
# Compatible for different versions of gym
env = CompatWrapper(env)
env = ActionMappingWrapper(env)
env.seed(args.seed)
obs_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
# Initialize model, algorithm, agent, replay_memory
model = MujocoModel(obs_dim, action_dim)
algorithm = SAC(
model,
gamma=GAMMA,
tau=TAU,
alpha=args.alpha,
actor_lr=ACTOR_LR,
critic_lr=CRITIC_LR)
agent = MujocoAgent(algorithm)
rpm = ReplayMemory(
max_size=MEMORY_SIZE, obs_dim=obs_dim, act_dim=action_dim)
total_steps = 0
test_flag = 0
while total_steps < args.train_total_steps:
# Train episode
episode_reward, episode_steps = run_train_episode(agent, env, rpm)
total_steps += episode_steps
summary.add_scalar('train/episode_reward', episode_reward, total_steps)
logger.info('Total Steps: {} Reward: {}'.format(
total_steps, episode_reward))
# Evaluate episode
if (total_steps + 1) // args.test_every_steps >= test_flag:
while (total_steps + 1) // args.test_every_steps >= test_flag:
test_flag += 1
avg_reward = run_evaluate_episodes(agent, env, EVAL_EPISODES)
summary.add_scalar('eval/episode_reward', avg_reward, total_steps)
logger.info('Evaluation over: {} episodes, Reward: {}'.format(
EVAL_EPISODES, avg_reward))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--env", default="HalfCheetah-v4", help='Mujoco gym environment name')
parser.add_argument("--seed", default=0, type=int, help='Sets Gym seed')
parser.add_argument(
"--train_total_steps",
default=3e6,
type=int,
help='Max time steps to run environment')
parser.add_argument(
'--test_every_steps',
type=int,
default=int(5e3),
help='The step interval between two consecutive evaluations')
parser.add_argument(
"--alpha",
default=0.2,
type=float,
help=
'Determines the relative importance of entropy term against the reward'
)
args = parser.parse_args()
main()