def make_atari_env(task, seed, training_num, test_num, **kwargs):
"""Wrapper function for Atari env.
If EnvPool is installed, it will automatically switch to EnvPool's Atari env.
:return: a tuple of (single env, training envs, test envs).
"""
if envpool is not None:
if kwargs.get("scale", 0):
warnings.warn(
"EnvPool does not include ScaledFloatFrame wrapper, "
"please set `x = x / 255.0` inside CNN network's forward function."
)
# parameters convertion
train_envs = env = envpool.make_gym(
task.replace("NoFrameskip-v4", "-v5"),
num_envs=training_num,
seed=seed,
episodic_life=True,
reward_clip=True,
stack_num=kwargs.get("frame_stack", 4),
)
test_envs = envpool.make_gym(
task.replace("NoFrameskip-v4", "-v5"),
num_envs=training_num,
seed=seed,
episodic_life=False,
reward_clip=False,
stack_num=kwargs.get("frame_stack", 4),
)
else:
warnings.warn("Recommend using envpool (pip install envpool) "
"to run Atari games more efficiently.")
env = wrap_deepmind(task, **kwargs)
train_envs = ShmemVectorEnv([
lambda: wrap_deepmind(
task, episode_life=True, clip_rewards=True, **kwargs)
for _ in range(training_num)
])
test_envs = ShmemVectorEnv([
lambda: wrap_deepmind(
task, episode_life=False, clip_rewards=False, **kwargs)
for _ in range(test_num)
])
env.seed(seed)
train_envs.seed(seed)
test_envs.seed(seed)
return env, train_envs, test_envs
def test_venv_wrapper_envpool_gym_reset_return_info():
num_envs = 4
env = VectorEnvNormObs(
envpool.make_gym("Ant-v3",
num_envs=num_envs,
gym_reset_return_info=True))
obs, info = env.reset()
assert obs.shape[0] == num_envs
for _, v in info.items():
if not isinstance(v, dict):
assert v.shape[0] == num_envs
def test_collector_envpool_gym_reset_return_info():
envs = envpool.make_gym("Pendulum-v1",
num_envs=4,
gym_reset_return_info=True)
policy = MyPolicy(action_shape=(len(envs), 1))
c0 = Collector(policy,
envs,
VectorReplayBuffer(len(envs) * 10, len(envs)),
exploration_noise=True)
c0.collect(n_step=8)
env_ids = np.zeros(len(envs) * 10)
env_ids[[0, 1, 10, 11, 20, 21, 30, 31]] = [0, 0, 1, 1, 2, 2, 3, 3]
assert np.allclose(c0.buffer.info["env_id"], env_ids)
def test_sac_with_il(args=get_args()):
# if you want to use python vector env, please refer to other test scripts
train_envs = env = envpool.make_gym(
args.task, num_envs=args.training_num, seed=args.seed
)
test_envs = envpool.make_gym(args.task, num_envs=args.test_num, seed=args.seed)
args.state_shape = env.observation_space.shape or env.observation_space.n
args.action_shape = env.action_space.shape or env.action_space.n
args.max_action = env.action_space.high[0]
if args.reward_threshold is None:
default_reward_threshold = {"Pendulum-v0": -250, "Pendulum-v1": -250}
args.reward_threshold = default_reward_threshold.get(
args.task, env.spec.reward_threshold
)
# you can also use tianshou.env.SubprocVectorEnv
# seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# model
net = Net(args.state_shape, hidden_sizes=args.hidden_sizes, device=args.device)
actor = ActorProb(
net,
args.action_shape,
max_action=args.max_action,
device=args.device,
unbounded=True
).to(args.device)
actor_optim = torch.optim.Adam(actor.parameters(), lr=args.actor_lr)
net_c1 = Net(
args.state_shape,
args.action_shape,
hidden_sizes=args.hidden_sizes,
concat=True,
device=args.device
)
critic1 = Critic(net_c1, device=args.device).to(args.device)
critic1_optim = torch.optim.Adam(critic1.parameters(), lr=args.critic_lr)
net_c2 = Net(
args.state_shape,
args.action_shape,
hidden_sizes=args.hidden_sizes,
concat=True,
device=args.device
)
critic2 = Critic(net_c2, device=args.device).to(args.device)
critic2_optim = torch.optim.Adam(critic2.parameters(), lr=args.critic_lr)
if args.auto_alpha:
target_entropy = -np.prod(env.action_space.shape)
log_alpha = torch.zeros(1, requires_grad=True, device=args.device)
alpha_optim = torch.optim.Adam([log_alpha], lr=args.alpha_lr)
args.alpha = (target_entropy, log_alpha, alpha_optim)
policy = SACPolicy(
actor,
actor_optim,
critic1,
critic1_optim,
critic2,
critic2_optim,
tau=args.tau,
gamma=args.gamma,
alpha=args.alpha,
reward_normalization=args.rew_norm,
estimation_step=args.n_step,
action_space=env.action_space
)
# collector
train_collector = Collector(
policy,
train_envs,
VectorReplayBuffer(args.buffer_size, len(train_envs)),
exploration_noise=True
)
test_collector = Collector(policy, test_envs)
# train_collector.collect(n_step=args.buffer_size)
# log
log_path = os.path.join(args.logdir, args.task, 'sac')
writer = SummaryWriter(log_path)
logger = TensorboardLogger(writer)
def save_best_fn(policy):
torch.save(policy.state_dict(), os.path.join(log_path, 'policy.pth'))
def stop_fn(mean_rewards):
return mean_rewards >= args.reward_threshold
# trainer
result = offpolicy_trainer(
policy,
train_collector,
test_collector,
args.epoch,
args.step_per_epoch,
args.step_per_collect,
args.test_num,
args.batch_size,
update_per_step=args.update_per_step,
stop_fn=stop_fn,
save_best_fn=save_best_fn,
logger=logger
)
assert stop_fn(result['best_reward'])
# here we define an imitation collector with a trivial policy
policy.eval()
if args.task.startswith("Pendulum"):
args.reward_threshold -= 50 # lower the goal
net = Actor(
Net(
args.state_shape,
hidden_sizes=args.imitation_hidden_sizes,
device=args.device
),
args.action_shape,
max_action=args.max_action,
device=args.device
).to(args.device)
optim = torch.optim.Adam(net.parameters(), lr=args.il_lr)
il_policy = ImitationPolicy(
net,
optim,
action_space=env.action_space,
action_scaling=True,
action_bound_method="clip"
)
il_test_collector = Collector(
il_policy,
envpool.make_gym(args.task, num_envs=args.test_num, seed=args.seed),
)
train_collector.reset()
result = offpolicy_trainer(
il_policy,
train_collector,
il_test_collector,
args.epoch,
args.il_step_per_epoch,
args.step_per_collect,
args.test_num,
args.batch_size,
stop_fn=stop_fn,
save_best_fn=save_best_fn,
logger=logger
)
assert stop_fn(result['best_reward'])
def test_psrl(args=get_args()):
train_envs = env = envpool.make_gym(args.task,
num_envs=args.training_num,
seed=args.seed)
test_envs = envpool.make_gym(args.task,
num_envs=args.test_num,
seed=args.seed)
if args.reward_threshold is None:
default_reward_threshold = {"NChain-v0": 3400}
args.reward_threshold = default_reward_threshold.get(
args.task, env.spec.reward_threshold)
print("reward threshold:", args.reward_threshold)
args.state_shape = env.observation_space.shape or env.observation_space.n
args.action_shape = env.action_space.shape or env.action_space.n
# seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# model
n_action = args.action_shape
n_state = args.state_shape
trans_count_prior = np.ones((n_state, n_action, n_state))
rew_mean_prior = np.full((n_state, n_action), args.rew_mean_prior)
rew_std_prior = np.full((n_state, n_action), args.rew_std_prior)
policy = PSRLPolicy(trans_count_prior, rew_mean_prior, rew_std_prior,
args.gamma, args.eps, args.add_done_loop)
# collector
train_collector = Collector(policy,
train_envs,
VectorReplayBuffer(args.buffer_size,
len(train_envs)),
exploration_noise=True)
test_collector = Collector(policy, test_envs)
# Logger
if args.logger == "wandb":
logger = WandbLogger(save_interval=1,
project='psrl',
name='wandb_test',
config=args)
if args.logger != "none":
log_path = os.path.join(args.logdir, args.task, 'psrl')
writer = SummaryWriter(log_path)
writer.add_text("args", str(args))
if args.logger == "tensorboard":
logger = TensorboardLogger(writer)
else:
logger.load(writer)
else:
logger = LazyLogger()
def stop_fn(mean_rewards):
return mean_rewards >= args.reward_threshold
train_collector.collect(n_step=args.buffer_size, random=True)
# trainer, test it without logger
result = onpolicy_trainer(
policy,
train_collector,
test_collector,
args.epoch,
args.step_per_epoch,
1,
args.test_num,
0,
episode_per_collect=args.episode_per_collect,
stop_fn=stop_fn,
logger=logger,
test_in_train=False,
)
if __name__ == '__main__':
pprint.pprint(result)
# Let's watch its performance!
policy.eval()
test_envs.seed(args.seed)
test_collector.reset()
result = test_collector.collect(n_episode=args.test_num,
render=args.render)
rews, lens = result["rews"], result["lens"]
print(f"Final reward: {rews.mean()}, length: {lens.mean()}")
elif env.spec.reward_threshold:
assert result["best_reward"] >= env.spec.reward_threshold
def test_a2c_with_il(args=get_args()):
# if you want to use python vector env, please refer to other test scripts
train_envs = env = envpool.make_gym(args.task,
num_envs=args.training_num,
seed=args.seed)
test_envs = envpool.make_gym(args.task,
num_envs=args.test_num,
seed=args.seed)
args.state_shape = env.observation_space.shape or env.observation_space.n
args.action_shape = env.action_space.shape or env.action_space.n
if args.reward_threshold is None:
default_reward_threshold = {"CartPole-v0": 195}
args.reward_threshold = default_reward_threshold.get(
args.task, env.spec.reward_threshold)
# seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# model
net = Net(args.state_shape,
hidden_sizes=args.hidden_sizes,
device=args.device)
actor = Actor(net, args.action_shape, device=args.device).to(args.device)
critic = Critic(net, device=args.device).to(args.device)
optim = torch.optim.Adam(ActorCritic(actor, critic).parameters(),
lr=args.lr)
dist = torch.distributions.Categorical
policy = A2CPolicy(actor,
critic,
optim,
dist,
discount_factor=args.gamma,
gae_lambda=args.gae_lambda,
vf_coef=args.vf_coef,
ent_coef=args.ent_coef,
max_grad_norm=args.max_grad_norm,
reward_normalization=args.rew_norm,
action_space=env.action_space)
# collector
train_collector = Collector(
policy, train_envs,
VectorReplayBuffer(args.buffer_size, len(train_envs)))
test_collector = Collector(policy, test_envs)
# log
log_path = os.path.join(args.logdir, args.task, 'a2c')
writer = SummaryWriter(log_path)
logger = TensorboardLogger(writer)
def save_best_fn(policy):
torch.save(policy.state_dict(), os.path.join(log_path, 'policy.pth'))
def stop_fn(mean_rewards):
return mean_rewards >= args.reward_threshold
# trainer
result = onpolicy_trainer(policy,
train_collector,
test_collector,
args.epoch,
args.step_per_epoch,
args.repeat_per_collect,
args.test_num,
args.batch_size,
episode_per_collect=args.episode_per_collect,
stop_fn=stop_fn,
save_best_fn=save_best_fn,
logger=logger)
assert stop_fn(result['best_reward'])
if __name__ == '__main__':
pprint.pprint(result)
# Let's watch its performance!
env = gym.make(args.task)
policy.eval()
collector = Collector(policy, env)
result = collector.collect(n_episode=1, render=args.render)
rews, lens = result["rews"], result["lens"]
print(f"Final reward: {rews.mean()}, length: {lens.mean()}")
policy.eval()
# here we define an imitation collector with a trivial policy
# if args.task == 'CartPole-v0':
# env.spec.reward_threshold = 190 # lower the goal
net = Net(args.state_shape,
hidden_sizes=args.hidden_sizes,
device=args.device)
net = Actor(net, args.action_shape, device=args.device).to(args.device)
optim = torch.optim.Adam(net.parameters(), lr=args.il_lr)
il_policy = ImitationPolicy(net, optim, action_space=env.action_space)
il_test_collector = Collector(
il_policy,
envpool.make_gym(args.task, num_envs=args.test_num, seed=args.seed),
)
train_collector.reset()
result = offpolicy_trainer(il_policy,
train_collector,
il_test_collector,
args.epoch,
args.il_step_per_epoch,
args.step_per_collect,
args.test_num,
args.batch_size,
stop_fn=stop_fn,
save_best_fn=save_best_fn,
logger=logger)
assert stop_fn(result['best_reward'])
if __name__ == '__main__':
pprint.pprint(result)
# Let's watch its performance!
env = gym.make(args.task)
il_policy.eval()
collector = Collector(il_policy, env)
result = collector.collect(n_episode=1, render=args.render)
rews, lens = result["rews"], result["lens"]
print(f"Final reward: {rews.mean()}, length: {lens.mean()}")