diff --git a/examples/pybullet/gym/pybullet_envs/stable_baselines/enjoy.py b/examples/pybullet/gym/pybullet_envs/stable_baselines/enjoy.py index 6c429abf7..1856d82bf 100644 --- a/examples/pybullet/gym/pybullet_envs/stable_baselines/enjoy.py +++ b/examples/pybullet/gym/pybullet_envs/stable_baselines/enjoy.py @@ -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() \ No newline at end of file + env.close() diff --git a/examples/pybullet/gym/pybullet_envs/stable_baselines/train.py b/examples/pybullet/gym/pybullet_envs/stable_baselines/train.py index 260cbbf9d..1eacc59b4 100644 --- a/examples/pybullet/gym/pybullet_envs/stable_baselines/train.py +++ b/examples/pybullet/gym/pybullet_envs/stable_baselines/train.py @@ -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) diff --git a/examples/pybullet/gym/pybullet_envs/stable_baselines/utils.py b/examples/pybullet/gym/pybullet_envs/stable_baselines/utils.py deleted file mode 100644 index fc64da7e7..000000000 --- a/examples/pybullet/gym/pybullet_envs/stable_baselines/utils.py +++ /dev/null @@ -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]))