mirror of
https://github.com/bulletphysics/bullet3
synced 2025-01-18 21:10:05 +00:00
84 lines
3.0 KiB
Python
84 lines
3.0 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
|
||
# You can run it using: python -m pybullet_envs.stable_baselines.enjoy --algo td3 --env HalfCheetahBulletEnv-v0
|
||
# Author: Antonin RAFFIN
|
||
# MIT License
|
||
import argparse
|
||
import multiprocessing
|
||
import time
|
||
import gym
|
||
import numpy as np
|
||
import pybullet_envs
|
||
|
||
from stable_baselines import SAC, TD3
|
||
|
||
from stable_baselines.common.evaluation import evaluate_policy
|
||
|
||
from pybullet_envs.stable_baselines.utils import TimeFeatureWrapper
|
||
|
||
|
||
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')
|
||
args = parser.parse_args()
|
||
|
||
env_id = args.env
|
||
# Create an env similar to the training env
|
||
env = TimeFeatureWrapper(gym.make(env_id))
|
||
|
||
# Use SubprocVecEnv for rendering
|
||
if not args.no_render:
|
||
env.render(mode='human')
|
||
|
||
algo = {
|
||
'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)
|
||
# Load the saved model
|
||
model = algo.load(save_path, env=env)
|
||
|
||
try:
|
||
# Use deterministic actions for evaluation
|
||
episode_rewards, episode_lengths = [], []
|
||
for _ in range(args.n_episodes):
|
||
obs = env.reset()
|
||
done = False
|
||
episode_reward = 0.0
|
||
episode_length = 0
|
||
while not done:
|
||
action, _ = model.predict(obs, deterministic=True)
|
||
obs, reward, done, _info = env.step(action)
|
||
episode_reward += reward
|
||
|
||
episode_length += 1
|
||
if not args.no_render:
|
||
env.render(mode='human')
|
||
dt = 1./240.
|
||
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))
|
||
|
||
mean_reward = np.mean(episode_rewards)
|
||
std_reward = np.std(episode_rewards)
|
||
|
||
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))
|
||
except KeyboardInterrupt:
|
||
pass
|
||
|
||
# Close process
|
||
env.close() |