522f7fbf98
This PR adds strict typing to the output of `update` and `learn` in all
policies. This will likely be the last large refactoring PR before the
next release (0.6.0, not 1.0.0), so it requires some attention. Several
difficulties were encountered on the path to that goal:
1. The policy hierarchy is actually "broken" in the sense that the keys
of dicts that were output by `learn` did not follow the same enhancement
(inheritance) pattern as the policies. This is a real problem and should
be addressed in the near future. Generally, several aspects of the
policy design and hierarchy might deserve a dedicated discussion.
2. Each policy needs to be generic in the stats return type, because one
might want to extend it at some point and then also extend the stats.
Even within the source code base this pattern is necessary in many
places.
3. The interaction between learn and update is a bit quirky, we
currently handle it by having update modify special field inside
TrainingStats, whereas all other fields are handled by learn.
4. The IQM module is a policy wrapper and required a
TrainingStatsWrapper. The latter relies on a bunch of black magic.
They were addressed by:
1. Live with the broken hierarchy, which is now made visible by bounds
in generics. We use type: ignore where appropriate.
2. Make all policies generic with bounds following the policy
inheritance hierarchy (which is incorrect, see above). We experimented a
bit with nested TrainingStats classes, but that seemed to add more
complexity and be harder to understand. Unfortunately, mypy thinks that
the code below is wrong, wherefore we have to add `type: ignore` to the
return of each `learn`
```python
T = TypeVar("T", bound=int)
def f() -> T:
return 3
```
3. See above
4. Write representative tests for the `TrainingStatsWrapper`. Still, the
black magic might cause nasty surprises down the line (I am not proud of
it)...
Closes #933
---------
Co-authored-by: Maximilian Huettenrauch <m.huettenrauch@appliedai.de>
Co-authored-by: Michael Panchenko <m.panchenko@appliedai.de>
104 lines
3.8 KiB
Python
104 lines
3.8 KiB
Python
import time
|
|
from collections.abc import Callable
|
|
from dataclasses import asdict
|
|
|
|
import numpy as np
|
|
|
|
from tianshou.data import (
|
|
Collector,
|
|
CollectStats,
|
|
InfoStats,
|
|
SequenceSummaryStats,
|
|
TimingStats,
|
|
)
|
|
from tianshou.policy import BasePolicy
|
|
from tianshou.utils import BaseLogger
|
|
|
|
|
|
def test_episode(
|
|
policy: BasePolicy,
|
|
collector: Collector,
|
|
test_fn: Callable[[int, int | None], None] | None,
|
|
epoch: int,
|
|
n_episode: int,
|
|
logger: BaseLogger | None = None,
|
|
global_step: int | None = None,
|
|
reward_metric: Callable[[np.ndarray], np.ndarray] | None = None,
|
|
) -> CollectStats:
|
|
"""A simple wrapper of testing policy in collector."""
|
|
collector.reset_env()
|
|
collector.reset_buffer()
|
|
policy.eval()
|
|
if test_fn:
|
|
test_fn(epoch, global_step)
|
|
result = collector.collect(n_episode=n_episode)
|
|
if reward_metric: # TODO: move into collector
|
|
rew = reward_metric(result.returns)
|
|
result.returns = rew
|
|
result.returns_stat = SequenceSummaryStats.from_sequence(rew)
|
|
if logger and global_step is not None:
|
|
assert result.n_collected_episodes > 0
|
|
logger.log_test_data(asdict(result), global_step)
|
|
return result
|
|
|
|
|
|
def gather_info(
|
|
start_time: float,
|
|
policy_update_time: float,
|
|
gradient_step: int,
|
|
best_reward: float,
|
|
best_reward_std: float,
|
|
train_collector: Collector | None = None,
|
|
test_collector: Collector | None = None,
|
|
) -> InfoStats:
|
|
"""A simple wrapper of gathering information from collectors.
|
|
|
|
:return: A dataclass object with the following members (depending on available collectors):
|
|
|
|
* ``gradient_step`` the total number of gradient steps;
|
|
* ``best_reward`` the best reward over the test results;
|
|
* ``best_reward_std`` the standard deviation of best reward over the test results;
|
|
* ``train_step`` the total collected step of training collector;
|
|
* ``train_episode`` the total collected episode of training collector;
|
|
* ``test_step`` the total collected step of test collector;
|
|
* ``test_episode`` the total collected episode of test collector;
|
|
* ``timing`` the timing statistics, with the following members:
|
|
* ``total_time`` the total time elapsed;
|
|
* ``train_time`` the total time elapsed for learning training (collecting samples plus model update);
|
|
* ``train_time_collect`` the time for collecting transitions in the \
|
|
training collector;
|
|
* ``train_time_update`` the time for training models;
|
|
* ``test_time`` the time for testing;
|
|
* ``update_speed`` the speed of updating (env_step per second).
|
|
"""
|
|
duration = max(0.0, time.time() - start_time)
|
|
test_time = 0.0
|
|
update_speed = 0.0
|
|
train_time_collect = 0.0
|
|
if test_collector is not None:
|
|
test_time = test_collector.collect_time
|
|
|
|
if train_collector is not None:
|
|
train_time_collect = train_collector.collect_time
|
|
update_speed = train_collector.collect_step / (duration - test_time)
|
|
|
|
timing_stat = TimingStats(
|
|
total_time=duration,
|
|
train_time=duration - test_time,
|
|
train_time_collect=train_time_collect,
|
|
train_time_update=policy_update_time,
|
|
test_time=test_time,
|
|
update_speed=update_speed,
|
|
)
|
|
|
|
return InfoStats(
|
|
gradient_step=gradient_step,
|
|
best_reward=best_reward,
|
|
best_reward_std=best_reward_std,
|
|
train_step=train_collector.collect_step if train_collector is not None else 0,
|
|
train_episode=train_collector.collect_episode if train_collector is not None else 0,
|
|
test_step=test_collector.collect_step if test_collector is not None else 0,
|
|
test_episode=test_collector.collect_episode if test_collector is not None else 0,
|
|
timing=timing_stat,
|
|
)
|