61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
# Code adapted from https://github.com/araffin/rl-baselines-zoo
|
||
# it requires stable-baselines to be installed
|
||
# Colab Notebook: https://colab.research.google.com/drive/1nZkHO4QTYfAksm9ZTaZ5vXyC7szZxC3F
|
||
# Author: Antonin RAFFIN
|
||
# MIT License
|
||
import argparse
|
||
|
||
import pybullet_envs
|
||
|
||
import gym
|
||
import numpy as np
|
||
from stable_baselines import SAC, TD3
|
||
from stable_baselines.common.noise import NormalActionNoise
|
||
|
||
from utils import TimeFeatureWrapper, EvalCallback
|
||
|
||
|
||
if __name__ == '__main__':
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--algo', help='RL Algorithm (Soft Actor-Critic by default)', default='sac',
|
||
type=str, required=False, choices=['sac', 'td3'])
|
||
parser.add_argument('--env', type=str, default='HalfCheetahBulletEnv-v0', help='environment ID')
|
||
parser.add_argument('-n', '--n-timesteps', help='Number of training timesteps', default=int(1e6),
|
||
type=int)
|
||
args = parser.parse_args()
|
||
|
||
env_id = args.env
|
||
n_timesteps = args.n_timesteps
|
||
save_path = '{}_{}'.format(args.algo, env_id)
|
||
|
||
# Instantiate and wrap the environment
|
||
env = TimeFeatureWrapper(gym.make(env_id))
|
||
|
||
# Create the evaluation environment and callback
|
||
eval_env = TimeFeatureWrapper(gym.make(env_id))
|
||
callback = EvalCallback(eval_env, best_model_save_path=save_path + '_best')
|
||
|
||
algo = {
|
||
'sac': SAC,
|
||
'td3': TD3
|
||
}[args.algo]
|
||
|
||
n_actions = env.action_space.shape[0]
|
||
|
||
# Tuned hyperparameters from https://github.com/araffin/rl-baselines-zoo
|
||
hyperparams = {
|
||
'sac': dict(batch_size=256, gamma=0.98, policy_kwargs=dict(layers=[256, 256]),
|
||
learning_starts=10000, buffer_size=int(2e5), tau=0.01),
|
||
|
||
'td3': dict(batch_size=100, policy_kwargs=dict(layers=[400, 300]),
|
||
learning_rate=1e-3, learning_starts=10000, buffer_size=int(1e6),
|
||
train_freq=1000, gradient_steps=1000,
|
||
action_noise=NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)))
|
||
}[args.algo]
|
||
|
||
model = algo('MlpPolicy', env, verbose=1, **hyperparams)
|
||
model.learn(n_timesteps, callback=callback)
|
||
|
||
print("Saving to {}.zip".format(save_path))
|
||
model.save(save_path)
|