def __init__(
self,
policy: BasePolicy,
train_collector: Collector,
test_collector: Optional[Collector],
max_epoch: int,
step_per_epoch: int,
repeat_per_collect: int,
episode_per_test: int,
batch_size: int,
step_per_collect: Optional[int] = None,
episode_per_collect: Optional[int] = None,
train_fn: Optional[Callable[[int, int], None]] = None,
test_fn: Optional[Callable[[int, Optional[int]], None]] = None,
stop_fn: Optional[Callable[[float], bool]] = None,
save_best_fn: Optional[Callable[[BasePolicy], None]] = None,
save_checkpoint_fn: Optional[Callable[[int, int, int], None]] = None,
resume_from_log: bool = False,
reward_metric: Optional[Callable[[np.ndarray], np.ndarray]] = None,
logger: BaseLogger = LazyLogger(),
verbose: bool = True,
show_progress: bool = True,
test_in_train: bool = True,
**kwargs: Any,
):
super().__init__(
learning_type="onpolicy",
policy=policy,
train_collector=train_collector,
test_collector=test_collector,
max_epoch=max_epoch,
step_per_epoch=step_per_epoch,
repeat_per_collect=repeat_per_collect,
episode_per_test=episode_per_test,
batch_size=batch_size,
step_per_collect=step_per_collect,
episode_per_collect=episode_per_collect,
train_fn=train_fn,
test_fn=test_fn,
stop_fn=stop_fn,
save_best_fn=save_best_fn,
save_checkpoint_fn=save_checkpoint_fn,
resume_from_log=resume_from_log,
reward_metric=reward_metric,
logger=logger,
verbose=verbose,
show_progress=show_progress,
test_in_train=test_in_train,
**kwargs,
)
def __init__(
self,
policy: BasePolicy,
buffer: ReplayBuffer,
test_collector: Optional[Collector],
max_epoch: int,
update_per_epoch: int,
episode_per_test: int,
batch_size: int,
test_fn: Optional[Callable[[int, Optional[int]], None]] = None,
stop_fn: Optional[Callable[[float], bool]] = None,
save_best_fn: Optional[Callable[[BasePolicy], None]] = None,
save_checkpoint_fn: Optional[Callable[[int, int, int], None]] = None,
resume_from_log: bool = False,
reward_metric: Optional[Callable[[np.ndarray], np.ndarray]] = None,
logger: BaseLogger = LazyLogger(),
verbose: bool = True,
**kwargs: Any,
):
super().__init__(
learning_type="offline",
policy=policy,
buffer=buffer,
test_collector=test_collector,
max_epoch=max_epoch,
update_per_epoch=update_per_epoch,
step_per_epoch=update_per_epoch,
episode_per_test=episode_per_test,
batch_size=batch_size,
test_fn=test_fn,
stop_fn=stop_fn,
save_best_fn=save_best_fn,
save_checkpoint_fn=save_checkpoint_fn,
resume_from_log=resume_from_log,
reward_metric=reward_metric,
logger=logger,
verbose=verbose,
**kwargs,
)
def onpolicy_trainer(
policy: BasePolicy,
train_collector: Collector,
test_collector: Collector,
max_epoch: int,
step_per_epoch: int,
repeat_per_collect: int,
episode_per_test: int,
batch_size: int,
step_per_collect: Optional[int] = None,
episode_per_collect: Optional[int] = None,
train_fn: Optional[Callable[[int, int], None]] = None,
test_fn: Optional[Callable[[int, Optional[int]], None]] = None,
stop_fn: Optional[Callable[[float], bool]] = None,
save_fn: Optional[Callable[[BasePolicy], None]] = None,
reward_metric: Optional[Callable[[np.ndarray], np.ndarray]] = None,
logger: BaseLogger = LazyLogger(),
verbose: bool = True,
test_in_train: bool = True,
) -> Dict[str, Union[float, str]]:
"""A wrapper for on-policy trainer procedure.
The "step" in trainer means an environment step (a.k.a. transition).
:param policy: an instance of the :class:`~tianshou.policy.BasePolicy` class.
:param Collector train_collector: the collector used for training.
:param Collector test_collector: the collector used for testing.
:param int max_epoch: the maximum number of epochs for training. The training
process might be finished before reaching ``max_epoch`` if ``stop_fn`` is set.
:param int step_per_epoch: the number of transitions collected per epoch.
:param int repeat_per_collect: the number of repeat time for policy learning, for
example, set it to 2 means the policy needs to learn each given batch data
twice.
:param int episode_per_test: the number of episodes for one policy evaluation.
:param int batch_size: the batch size of sample data, which is going to feed in the
policy network.
:param int step_per_collect: the number of transitions the collector would collect
before the network update, i.e., trainer will collect "step_per_collect"
transitions and do some policy network update repeatly in each epoch.
:param int episode_per_collect: the number of episodes the collector would collect
before the network update, i.e., trainer will collect "episode_per_collect"
episodes and do some policy network update repeatly in each epoch.
:param function train_fn: a hook called at the beginning of training in each epoch.
It can be used to perform custom additional operations, with the signature ``f(
num_epoch: int, step_idx: int) -> None``.
:param function test_fn: a hook called at the beginning of testing in each epoch.
It can be used to perform custom additional operations, with the signature ``f(
num_epoch: int, step_idx: int) -> None``.
:param function save_fn: a hook called when the undiscounted average mean reward in
evaluation phase gets better, with the signature ``f(policy: BasePolicy) ->
None``.
:param function stop_fn: a function with signature ``f(mean_rewards: float) ->
bool``, receives the average undiscounted returns of the testing result,
returns a boolean which indicates whether reaching the goal.
:param function reward_metric: a function with signature ``f(rewards: np.ndarray
with shape (num_episode, agent_num)) -> np.ndarray with shape (num_episode,)``,
used in multi-agent RL. We need to return a single scalar for each episode's
result to monitor training in the multi-agent RL setting. This function
specifies what is the desired metric, e.g., the reward of agent 1 or the
average reward over all agents.
:param BaseLogger logger: A logger that logs statistics during
training/testing/updating. Default to a logger that doesn't log anything.
:param bool verbose: whether to print the information. Default to True.
:param bool test_in_train: whether to test in the training phase. Default to True.
:return: See :func:`~tianshou.trainer.gather_info`.
.. note::
Only either one of step_per_collect and episode_per_collect can be specified.
"""
env_step, gradient_step = 0, 0
last_rew, last_len = 0.0, 0
stat: Dict[str, MovAvg] = defaultdict(MovAvg)
start_time = time.time()
train_collector.reset_stat()
test_collector.reset_stat()
test_in_train = test_in_train and train_collector.policy == policy
test_result = test_episode(policy, test_collector, test_fn, 0,
episode_per_test, logger, env_step,
reward_metric)
best_epoch = 0
best_reward, best_reward_std = test_result["rew"], test_result["rew_std"]
for epoch in range(1, 1 + max_epoch):
# train
policy.train()
with tqdm.tqdm(total=step_per_epoch,
desc=f"Epoch #{epoch}",
**tqdm_config) as t:
while t.n < t.total:
if train_fn:
train_fn(epoch, env_step)
result = train_collector.collect(n_step=step_per_collect,
n_episode=episode_per_collect)
if reward_metric:
result["rews"] = reward_metric(result["rews"])
env_step += int(result["n/st"])
t.update(result["n/st"])
logger.log_train_data(result, env_step)
last_rew = result['rew'] if 'rew' in result else last_rew
last_len = result['len'] if 'len' in result else last_len
data = {
"env_step": str(env_step),
"rew": f"{last_rew:.2f}",
"len": str(int(last_len)),
"n/ep": str(int(result["n/ep"])),
"n/st": str(int(result["n/st"])),
}
if test_in_train and stop_fn and stop_fn(result["rew"]):
test_result = test_episode(policy, test_collector, test_fn,
epoch, episode_per_test, logger,
env_step)
if stop_fn(test_result["rew"]):
if save_fn:
save_fn(policy)
t.set_postfix(**data)
return gather_info(start_time, train_collector,
test_collector, test_result["rew"],
test_result["rew_std"])
else:
policy.train()
losses = policy.update(0,
train_collector.buffer,
batch_size=batch_size,
repeat=repeat_per_collect)
train_collector.reset_buffer()
step = max(
[1] +
[len(v) for v in losses.values() if isinstance(v, list)])
gradient_step += step
for k in losses.keys():
stat[k].add(losses[k])
losses[k] = stat[k].get()
data[k] = f"{losses[k]:.3f}"
logger.log_update_data(losses, gradient_step)
t.set_postfix(**data)
if t.n <= t.total:
t.update()
# test
test_result = test_episode(policy, test_collector, test_fn, epoch,
episode_per_test, logger, env_step)
rew, rew_std = test_result["rew"], test_result["rew_std"]
if best_epoch == -1 or best_reward < rew:
best_reward, best_reward_std = rew, rew_std
best_epoch = epoch
if save_fn:
save_fn(policy)
if verbose:
print(
f"Epoch #{epoch}: test_reward: {rew:.6f} ± {rew_std:.6f}, best_rew"
f"ard: {best_reward:.6f} ± {best_reward_std:.6f} in #{best_epoch}"
)
if stop_fn and stop_fn(best_reward):
break
return gather_info(start_time, train_collector, test_collector,
best_reward, best_reward_std)
def offline_trainer(
policy: BasePolicy,
buffer: ReplayBuffer,
test_collector: Collector,
max_epoch: int,
update_per_epoch: int,
episode_per_test: int,
batch_size: int,
test_fn: Optional[Callable[[int, Optional[int]], None]] = None,
stop_fn: Optional[Callable[[float], bool]] = None,
save_fn: Optional[Callable[[BasePolicy], None]] = None,
reward_metric: Optional[Callable[[np.ndarray], np.ndarray]] = None,
logger: BaseLogger = LazyLogger(),
verbose: bool = True,
) -> Dict[str, Union[float, str]]:
"""A wrapper for offline trainer procedure.
The "step" in offline trainer means a gradient step.
:param policy: an instance of the :class:`~tianshou.policy.BasePolicy` class.
:param Collector test_collector: the collector used for testing.
:param int max_epoch: the maximum number of epochs for training. The training
process might be finished before reaching ``max_epoch`` if ``stop_fn`` is set.
:param int update_per_epoch: the number of policy network updates, so-called
gradient steps, per epoch.
:param episode_per_test: the number of episodes for one policy evaluation.
:param int batch_size: the batch size of sample data, which is going to feed in
the policy network.
:param function test_fn: a hook called at the beginning of testing in each epoch.
It can be used to perform custom additional operations, with the signature ``f(
num_epoch: int, step_idx: int) -> None``.
:param function save_fn: a hook called when the undiscounted average mean reward in
evaluation phase gets better, with the signature ``f(policy: BasePolicy) ->
None``.
:param function stop_fn: a function with signature ``f(mean_rewards: float) ->
bool``, receives the average undiscounted returns of the testing result,
returns a boolean which indicates whether reaching the goal.
:param function reward_metric: a function with signature ``f(rewards: np.ndarray
with shape (num_episode, agent_num)) -> np.ndarray with shape (num_episode,)``,
used in multi-agent RL. We need to return a single scalar for each episode's
result to monitor training in the multi-agent RL setting. This function
specifies what is the desired metric, e.g., the reward of agent 1 or the
average reward over all agents.
:param BaseLogger logger: A logger that logs statistics during updating/testing.
Default to a logger that doesn't log anything.
:param bool verbose: whether to print the information. Default to True.
:return: See :func:`~tianshou.trainer.gather_info`.
"""
gradient_step = 0
stat: Dict[str, MovAvg] = defaultdict(MovAvg)
start_time = time.time()
test_collector.reset_stat()
test_result = test_episode(policy, test_collector, test_fn, 0,
episode_per_test, logger, gradient_step,
reward_metric)
best_epoch = 0
best_reward, best_reward_std = test_result["rew"], test_result["rew_std"]
for epoch in range(1, 1 + max_epoch):
policy.train()
with tqdm.trange(update_per_epoch,
desc=f"Epoch #{epoch}",
**tqdm_config) as t:
for i in t:
gradient_step += 1
losses = policy.update(batch_size, buffer)
data = {"gradient_step": str(gradient_step)}
for k in losses.keys():
stat[k].add(losses[k])
losses[k] = stat[k].get()
data[k] = f"{losses[k]:.3f}"
logger.log_update_data(losses, gradient_step)
t.set_postfix(**data)
# test
test_result = test_episode(policy, test_collector, test_fn, epoch,
episode_per_test, logger, gradient_step,
reward_metric)
rew, rew_std = test_result["rew"], test_result["rew_std"]
if best_epoch == -1 or best_reward < rew:
best_reward, best_reward_std = rew, rew_std
best_epoch = epoch
if save_fn:
save_fn(policy)
if verbose:
print(
f"Epoch #{epoch}: test_reward: {rew:.6f} ± {rew_std:.6f}, best_rew"
f"ard: {best_reward:.6f} ± {best_reward_std:.6f} in #{best_epoch}"
)
if stop_fn and stop_fn(best_reward):
break
return gather_info(start_time, None, test_collector, best_reward,
best_reward_std)
def __init__(
self,
learning_type: str,
policy: BasePolicy,
max_epoch: int,
batch_size: int,
train_collector: Optional[Collector] = None,
test_collector: Optional[Collector] = None,
buffer: Optional[ReplayBuffer] = None,
step_per_epoch: Optional[int] = None,
repeat_per_collect: Optional[int] = None,
episode_per_test: Optional[int] = None,
update_per_step: Union[int, float] = 1,
update_per_epoch: Optional[int] = None,
step_per_collect: Optional[int] = None,
episode_per_collect: Optional[int] = None,
train_fn: Optional[Callable[[int, int], None]] = None,
test_fn: Optional[Callable[[int, Optional[int]], None]] = None,
stop_fn: Optional[Callable[[float], bool]] = None,
save_best_fn: Optional[Callable[[BasePolicy], None]] = None,
save_checkpoint_fn: Optional[Callable[[int, int, int], str]] = None,
resume_from_log: bool = False,
reward_metric: Optional[Callable[[np.ndarray], np.ndarray]] = None,
logger: BaseLogger = LazyLogger(),
verbose: bool = True,
show_progress: bool = True,
test_in_train: bool = True,
save_fn: Optional[Callable[[BasePolicy], None]] = None,
):
if save_fn:
deprecation(
"save_fn in trainer is marked as deprecated and will be "
"removed in the future. Please use save_best_fn instead.")
assert save_best_fn is None
save_best_fn = save_fn
self.policy = policy
self.buffer = buffer
self.train_collector = train_collector
self.test_collector = test_collector
self.logger = logger
self.start_time = time.time()
self.stat: DefaultDict[str, MovAvg] = defaultdict(MovAvg)
self.best_reward = 0.0
self.best_reward_std = 0.0
self.start_epoch = 0
self.gradient_step = 0
self.env_step = 0
self.max_epoch = max_epoch
self.step_per_epoch = step_per_epoch
# either on of these two
self.step_per_collect = step_per_collect
self.episode_per_collect = episode_per_collect
self.update_per_step = update_per_step
self.repeat_per_collect = repeat_per_collect
self.episode_per_test = episode_per_test
self.batch_size = batch_size
self.train_fn = train_fn
self.test_fn = test_fn
self.stop_fn = stop_fn
self.save_best_fn = save_best_fn
self.save_checkpoint_fn = save_checkpoint_fn
self.reward_metric = reward_metric
self.verbose = verbose
self.show_progress = show_progress
self.test_in_train = test_in_train
self.resume_from_log = resume_from_log
self.is_run = False
self.last_rew, self.last_len = 0.0, 0
self.epoch = self.start_epoch
self.best_epoch = self.start_epoch
self.stop_fn_flag = False
self.iter_num = 0
def test_psrl(args=get_args()):
env = gym.make(args.task)
if args.task == "NChain-v0":
env.spec.reward_threshold = 3400
# env.spec.reward_threshold = 3647 # described in PSRL paper
print("reward threshold:", env.spec.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
# train_envs = gym.make(args.task)
train_envs = DummyVectorEnv(
[lambda: gym.make(args.task) for _ in range(args.training_num)])
# test_envs = gym.make(args.task)
test_envs = SubprocVectorEnv(
[lambda: gym.make(args.task) for _ in range(args.test_num)])
# seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
train_envs.seed(args.seed)
test_envs.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)
elif args.logger == "tensorboard":
log_path = os.path.join(args.logdir, args.task, 'psrl')
writer = SummaryWriter(log_path)
writer.add_text("args", str(args))
logger = TensorboardLogger(writer)
else:
logger = LazyLogger()
def stop_fn(mean_rewards):
if env.spec.reward_threshold:
return mean_rewards >= env.spec.reward_threshold
else:
return False
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 offpolicy_trainer(
policy: BasePolicy,
train_collector: Collector,
test_collector: Collector,
max_epoch: int,
step_per_epoch: int,
step_per_collect: int,
episode_per_test: int,
batch_size: int,
update_per_step: Union[int, float] = 1,
train_fn: Optional[Callable[[int, int], None]] = None,
test_fn: Optional[Callable[[int, Optional[int]], None]] = None,
stop_fn: Optional[Callable[[float], bool]] = None,
save_fn: Optional[Callable[[BasePolicy], None]] = None,
save_checkpoint_fn: Optional[Callable[[int, int, int], None]] = None,
resume_from_log: bool = False,
reward_metric: Optional[Callable[[np.ndarray], np.ndarray]] = None,
logger: BaseLogger = LazyLogger(),
verbose: bool = True,
test_in_train: bool = True,
) -> Dict[str, Union[float, str]]:
"""A wrapper for off-policy trainer procedure.
The "step" in trainer means an environment step (a.k.a. transition).
:param policy: an instance of the :class:`~tianshou.policy.BasePolicy` class.
:param Collector train_collector: the collector used for training.
:param Collector test_collector: the collector used for testing.
:param int max_epoch: the maximum number of epochs for training. The training
process might be finished before reaching ``max_epoch`` if ``stop_fn`` is set.
:param int step_per_epoch: the number of transitions collected per epoch.
:param int step_per_collect: the number of transitions the collector would collect
before the network update, i.e., trainer will collect "step_per_collect"
transitions and do some policy network update repeatly in each epoch.
:param episode_per_test: the number of episodes for one policy evaluation.
:param int batch_size: the batch size of sample data, which is going to feed in the
policy network.
:param int/float update_per_step: the number of times the policy network would be
updated per transition after (step_per_collect) transitions are collected,
e.g., if update_per_step set to 0.3, and step_per_collect is 256, policy will
be updated round(256 * 0.3 = 76.8) = 77 times after 256 transitions are
collected by the collector. Default to 1.
:param function train_fn: a hook called at the beginning of training in each epoch.
It can be used to perform custom additional operations, with the signature ``f(
num_epoch: int, step_idx: int) -> None``.
:param function test_fn: a hook called at the beginning of testing in each epoch.
It can be used to perform custom additional operations, with the signature ``f(
num_epoch: int, step_idx: int) -> None``.
:param function save_fn: a hook called when the undiscounted average mean reward in
evaluation phase gets better, with the signature ``f(policy: BasePolicy) ->
None``.
:param function save_checkpoint_fn: a function to save training process, with the
signature ``f(epoch: int, env_step: int, gradient_step: int) -> None``; you can
save whatever you want.
:param bool resume_from_log: resume env_step/gradient_step and other metadata from
existing tensorboard log. Default to False.
:param function stop_fn: a function with signature ``f(mean_rewards: float) ->
bool``, receives the average undiscounted returns of the testing result,
returns a boolean which indicates whether reaching the goal.
:param function reward_metric: a function with signature ``f(rewards: np.ndarray
with shape (num_episode, agent_num)) -> np.ndarray with shape (num_episode,)``,
used in multi-agent RL. We need to return a single scalar for each episode's
result to monitor training in the multi-agent RL setting. This function
specifies what is the desired metric, e.g., the reward of agent 1 or the
average reward over all agents.
:param BaseLogger logger: A logger that logs statistics during
training/testing/updating. Default to a logger that doesn't log anything.
:param bool verbose: whether to print the information. Default to True.
:param bool test_in_train: whether to test in the training phase. Default to True.
:return: See :func:`~tianshou.trainer.gather_info`.
"""
if save_fn:
warnings.warn("Please consider using save_checkpoint_fn instead of save_fn.")
start_epoch, env_step, gradient_step = 0, 0, 0
if resume_from_log:
start_epoch, env_step, gradient_step = logger.restore_data()
last_rew, last_len = 0.0, 0
stat: Dict[str, MovAvg] = defaultdict(MovAvg)
start_time = time.time()
train_collector.reset_stat()
test_collector.reset_stat()
test_in_train = test_in_train and train_collector.policy == policy
test_result = test_episode(policy, test_collector, test_fn, start_epoch,
episode_per_test, logger, env_step, reward_metric)
best_epoch = start_epoch
best_reward, best_reward_std = test_result["rew"], test_result["rew_std"]
for epoch in range(1 + start_epoch, 1 + max_epoch):
# train
policy.train()
with tqdm.tqdm(
total=step_per_epoch, desc=f"Epoch #{epoch}", **tqdm_config
) as t:
while t.n < t.total:
if train_fn:
train_fn(epoch, env_step)
result = train_collector.collect(n_step=step_per_collect)
if result["n/ep"] > 0 and reward_metric:
result["rews"] = reward_metric(result["rews"])
env_step += int(result["n/st"])
t.update(result["n/st"])
logger.log_train_data(result, env_step)
last_rew = result['rew'] if 'rew' in result else last_rew
last_len = result['len'] if 'len' in result else last_len
data = {
"env_step": str(env_step),
"rew": f"{last_rew:.2f}",
"len": str(int(last_len)),
"n/ep": str(int(result["n/ep"])),
"n/st": str(int(result["n/st"])),
}
if result["n/ep"] > 0:
if test_in_train and stop_fn and stop_fn(result["rew"]):
test_result = test_episode(
policy, test_collector, test_fn,
epoch, episode_per_test, logger, env_step)
if stop_fn(test_result["rew"]):
if save_fn:
save_fn(policy)
logger.save_data(
epoch, env_step, gradient_step, save_checkpoint_fn)
t.set_postfix(**data)
return gather_info(
start_time, train_collector, test_collector,
test_result["rew"], test_result["rew_std"])
else:
policy.train()
for i in range(round(update_per_step * result["n/st"])):
gradient_step += 1
losses = policy.update(batch_size, train_collector.buffer)
for k in losses.keys():
stat[k].add(losses[k])
losses[k] = stat[k].get()
data[k] = f"{losses[k]:.3f}"
logger.log_update_data(losses, gradient_step)
t.set_postfix(**data)
if t.n <= t.total:
t.update()
# test
test_result = test_episode(policy, test_collector, test_fn, epoch,
episode_per_test, logger, env_step, reward_metric)
rew, rew_std = test_result["rew"], test_result["rew_std"]
if best_epoch < 0 or best_reward < rew:
best_epoch, best_reward, best_reward_std = epoch, rew, rew_std
if save_fn:
save_fn(policy)
logger.save_data(epoch, env_step, gradient_step, save_checkpoint_fn)
if verbose:
print(f"Epoch #{epoch}: test_reward: {rew:.6f} ± {rew_std:.6f}, best_rew"
f"ard: {best_reward:.6f} ± {best_reward_std:.6f} in #{best_epoch}")
if stop_fn and stop_fn(best_reward):
break
return gather_info(start_time, train_collector, test_collector,
best_reward, best_reward_std)
def test_psrl(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)
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
