Scarfski/bullet3
1
0
Fork 0
mirror of https://github.com/bulletphysics/bullet3.git synced 2026-10-12 11:55:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Migrate to Stable-Baselines3

Browse Source
This commit is contained in:
Antonin Raffin
2021-07-29 15:21:49 +02:00
parent 0e124cb2f1
commit 2f8b44c702
3 changed files with 114 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
examples/pybullet/gym/pybullet_envs/stable_baselines/enjoy.py
View File

@@ -1,4 +1,4 @@
# Code adapted from https://github.com/araffin/rl-baselines-zoo
# 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
# You can run it using: python -m pybullet_envs.stable_baselines.enjoy --algo td3 --env HalfCheetahBulletEnv-v0
@@ -7,50 +7,66 @@
import os
import time
import argparse
import multiprocessing
import gym
import numpy as np
import pybullet_envs
from stable_baselines import SAC, TD3
from pybullet_envs.stable_baselines.utils import TimeFeatureWrapper
from stable_baselines3 import SAC, TD3
if __name__ == '__main__':
parser = argparse.ArgumentParser("Enjoy an RL agent trained using Stable Baselines")
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-episodes', help='Number of episodes', default=5,
type=int)
parser.add_argument('--no-render', action='store_true', default=False,
help='Do not render the environment')
parser.add_argument('--load-best', action='store_true', default=False,
help='Load best model instead of last model if available')
if __name__ == "__main__":
parser = argparse.ArgumentParser(
"Enjoy an RL agent trained using Stable Baselines3"
)
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-episodes", help="Number of episodes", default=5, type=int
)
parser.add_argument(
"--no-render",
action="store_true",
default=False,
help="Do not render the environment",
)
parser.add_argument(
"--load-best",
action="store_true",
default=False,
help="Load best model instead of last model if available",
)
args = parser.parse_args()
env_id = args.env
# Create an env similar to the training env
env = TimeFeatureWrapper(gym.make(env_id))
env = gym.make(env_id)
# Use SubprocVecEnv for rendering
# Enable GUI
if not args.no_render:
env.render(mode='human')
env.render(mode="human")
algo = {
'sac': SAC,
'td3': TD3
"sac": SAC,
"td3": TD3,
}[args.algo]
# We assume that the saved model is in the same folder
save_path = '{}_{}.zip'.format(args.algo, env_id)
save_path = f"{args.algo}_{env_id}.zip"
if not os.path.isfile(save_path) or args.load_best:
print("Loading best model")
# Try to load best model
save_path = os.path.join('{}_{}'.format(args.algo, env_id), 'best_model.zip')
save_path = os.path.join(f"{args.algo}_{env_id}", "best_model.zip")
# Load the saved model
model = algo.load(save_path, env=env)
@@ -70,12 +86,14 @@ if __name__ == '__main__':
episode_length += 1
if not args.no_render:
env.render(mode='human')
dt = 1. / 240.
env.render(mode="human")
dt = 1.0 / 240.0
time.sleep(dt)
episode_rewards.append(episode_reward)
episode_lengths.append(episode_length)
print("Episode {} reward={}, length={}".format(len(episode_rewards), episode_reward, episode_length))
print(
f"Episode {len(episode_rewards)} reward={episode_reward}, length={episode_length}"
)
mean_reward = np.mean(episode_rewards)
std_reward = np.std(episode_rewards)
@@ -83,10 +101,10 @@ if __name__ == '__main__':
mean_len, std_len = np.mean(episode_lengths), np.std(episode_lengths)
print("==== Results ====")
print("Episode_reward={:.2f} +/- {:.2f}".format(mean_reward, std_reward))
print("Episode_length={:.2f} +/- {:.2f}".format(mean_len, std_len))
print(f"Episode_reward={mean_reward:.2f} +/- {std_reward:.2f}")
print(f"Episode_length={mean_len:.2f} +/- {std_len:.2f}")
except KeyboardInterrupt:
pass
# Close process
env.close()
env.close()

103
examples/pybullet/gym/pybullet_envs/stable_baselines/train.py
View File

@@ -1,74 +1,105 @@
# 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
# Code adapted from https://github.com/DLR-RM/rl-baselines3-zoo
# it requires stable-baselines3 to be installed
# Colab Notebook: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/pybullet.ipynb
# You can run it using: python -m pybullet_envs.stable_baselines.train --algo td3 --env HalfCheetahBulletEnv-v0
# Author: Antonin RAFFIN
# MIT License
import argparse
import pybullet_envs
import pybullet_envs # register pybullet envs
import gym
import numpy as np
from stable_baselines import SAC, TD3
from stable_baselines.common.noise import NormalActionNoise
from stable_baselines.common.callbacks import EvalCallback, CheckpointCallback
from stable_baselines.common.vec_env import DummyVecEnv
from pybullet_envs.stable_baselines.utils import TimeFeatureWrapper
from stable_baselines3 import SAC, TD3
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import EvalCallback, CheckpointCallback
from stable_baselines3.common.monitor import Monitor
if __name__ == '__main__':
parser = argparse.ArgumentParser("Train an RL agent using Stable Baselines")
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)
parser.add_argument('--save-freq', help='Save the model every n steps (if negative, no checkpoint)',
default=-1, type=int)
if __name__ == "__main__":
parser = argparse.ArgumentParser("Train an RL agent using Stable Baselines3")
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,
)
parser.add_argument(
"--save-freq",
help="Save the model every n steps (if negative, no checkpoint)",
default=-1,
type=int,
)
args = parser.parse_args()
env_id = args.env
n_timesteps = args.n_timesteps
save_path = '{}_{}'.format(args.algo, env_id)
save_path = f"{args.algo}_{env_id}"
# Instantiate and wrap the environment
env = TimeFeatureWrapper(gym.make(env_id))
env = gym.make(env_id)
# Create the evaluation environment and callbacks
eval_env = DummyVecEnv([lambda: TimeFeatureWrapper(gym.make(env_id))])
eval_env = Monitor(gym.make(env_id))
callbacks = [EvalCallback(eval_env, best_model_save_path=save_path)]
# Save a checkpoint every n steps
if args.save_freq > 0:
callbacks.append(CheckpointCallback(save_freq=args.save_freq, save_path=save_path,
name_prefix='rl_model'))
callbacks.append(
CheckpointCallback(
save_freq=args.save_freq, save_path=save_path, name_prefix="rl_model"
)
)
algo = {
'sac': SAC,
'td3': TD3
"sac": SAC,
"td3": TD3,
}[args.algo]
n_actions = env.action_space.shape[0]
# Tuned hyperparameters from https://github.com/araffin/rl-baselines-zoo
# Tuned hyperparameters from https://github.com/DLR-RM/rl-baselines3-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)))
"sac": dict(
batch_size=256,
gamma=0.98,
policy_kwargs=dict(net_arch=[256, 256]),
learning_starts=10000,
buffer_size=int(3e5),
tau=0.01,
),
"td3": dict(
batch_size=100,
policy_kwargs=dict(net_arch=[400, 300]),
learning_rate=1e-3,
learning_starts=10000,
buffer_size=int(1e6),
train_freq=1,
gradient_steps=1,
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 = algo("MlpPolicy", env, verbose=1, **hyperparams)
try:
model.learn(n_timesteps, callback=callbacks)
except KeyboardInterrupt:
pass
print("Saving to {}.zip".format(save_path))
print(f"Saving to {save_path}.zip")
model.save(save_path)

59
examples/pybullet/gym/pybullet_envs/stable_baselines/utils.py
View File

@@ -1,59 +0,0 @@
# Code adapted from https://github.com/araffin/rl-baselines-zoo
# it requires stable-baselines to be installed
# Author: Antonin RAFFIN
# MIT License
import gym
import numpy as np
from gym.wrappers import TimeLimit
class TimeFeatureWrapper(gym.Wrapper):
"""
Add remaining time to observation space for fixed length episodes.
See https://arxiv.org/abs/1712.00378 and https://github.com/aravindr93/mjrl/issues/13.
:param env: (gym.Env)
:param max_steps: (int) Max number of steps of an episode
if it is not wrapped in a TimeLimit object.
:param test_mode: (bool) In test mode, the time feature is constant,
equal to zero. This allow to check that the agent did not overfit this feature,
learning a deterministic pre-defined sequence of actions.
"""
def __init__(self, env, max_steps=1000, test_mode=False):
assert isinstance(env.observation_space, gym.spaces.Box)
# Add a time feature to the observation
low, high = env.observation_space.low, env.observation_space.high
low, high = np.concatenate((low, [0])), np.concatenate((high, [1.]))
env.observation_space = gym.spaces.Box(low=low, high=high, dtype=np.float32)
super(TimeFeatureWrapper, self).__init__(env)
if isinstance(env, TimeLimit):
self._max_steps = env._max_episode_steps
else:
self._max_steps = max_steps
self._current_step = 0
self._test_mode = test_mode
def reset(self):
self._current_step = 0
return self._get_obs(self.env.reset())
def step(self, action):
self._current_step += 1
obs, reward, done, info = self.env.step(action)
return self._get_obs(obs), reward, done, info
def _get_obs(self, obs):
"""
Concatenate the time feature to the current observation.
:param obs: (np.ndarray)
:return: (np.ndarray)
"""
# Remaining time is more general
time_feature = 1 - (self._current_step / self._max_steps)
if self._test_mode:
time_feature = 1.0
# Optionnaly: concatenate [time_feature, time_feature ** 2]
return np.concatenate((obs, [time_feature]))