Michael Panchenko 07702fc007
Improved typing and reduced duplication (#912)
# Goals of the PR

The PR introduces **no changes to functionality**, apart from improved
input validation here and there. The main goals are to reduce some
complexity of the code, to improve types and IDE completions, and to
extend documentation and block comments where appropriate. Because of
the change to the trainer interfaces, many files are affected (more
details below), but still the overall changes are "small" in a certain
sense.

## Major Change 1 - BatchProtocol

**TL;DR:** One can now annotate which fields the batch is expected to
have on input params and which fields a returned batch has. Should be
useful for reading the code. getting meaningful IDE support, and
catching bugs with mypy. This annotation strategy will continue to work
if Batch is replaced by TensorDict or by something else.

**In more detail:** Batch itself has no fields and using it for
annotations is of limited informational power. Batches with fields are
not separate classes but instead instances of Batch directly, so there
is no type that could be used for annotation. Fortunately, python
`Protocol` is here for the rescue. With these changes we can now do
things like

```python
class ActionBatchProtocol(BatchProtocol):
    logits: Sequence[Union[tuple, torch.Tensor]]
    dist: torch.distributions.Distribution
    act: torch.Tensor
    state: Optional[torch.Tensor]


class RolloutBatchProtocol(BatchProtocol):
    obs: torch.Tensor
    obs_next: torch.Tensor
    info: Dict[str, Any]
    rew: torch.Tensor
    terminated: torch.Tensor
    truncated: torch.Tensor

class PGPolicy(BasePolicy):
    ...

    def forward(
        self,
        batch: RolloutBatchProtocol,
        state: Optional[Union[dict, Batch, np.ndarray]] = None,
        **kwargs: Any,
    ) -> ActionBatchProtocol:

```

The IDE and mypy are now very helpful in finding errors and in
auto-completion, whereas before the tools couldn't assist in that at
all.

## Major Change 2 - remove duplication in trainer package

**TL;DR:** There was a lot of duplication between `BaseTrainer` and its
subclasses. Even worse, it was almost-duplication. There was also
interface fragmentation through things like `onpolicy_trainer`. Now this
duplication is gone and all downstream code was adjusted.

**In more detail:** Since this change affects a lot of code, I would
like to explain why I thought it to be necessary.

1. The subclasses of `BaseTrainer` just duplicated docstrings and
constructors. What's worse, they changed the order of args there, even
turning some kwargs of BaseTrainer into args. They also had the arg
`learning_type` which was passed as kwarg to the base class and was
unused there. This made things difficult to maintain, and in fact some
errors were already present in the duplicated docstrings.
2. The "functions" a la `onpolicy_trainer`, which just called the
`OnpolicyTrainer.run`, not only introduced interface fragmentation but
also completely obfuscated the docstring and interfaces. They themselves
had no dosctring and the interface was just `*args, **kwargs`, which
makes it impossible to understand what they do and which things can be
passed without reading their implementation, then reading the docstring
of the associated class, etc. Needless to say, mypy and IDEs provide no
support with such functions. Nevertheless, they were used everywhere in
the code-base. I didn't find the sacrifices in clarity and complexity
justified just for the sake of not having to write `.run()` after
instantiating a trainer.
3. The trainers are all very similar to each other. As for my
application I needed a new trainer, I wanted to understand their
structure. The similarity, however, was hard to discover since they were
all in separate modules and there was so much duplication. I kept
staring at the constructors for a while until I figured out that
essentially no changes to the superclass were introduced. Now they are
all in the same module and the similarities/differences between them are
much easier to grasp (in my opinion)
4. Because of (1), I had to manually change and check a lot of code,
which was very tedious and boring. This kind of work won't be necessary
in the future, since now IDEs can be used for changing signatures,
renaming args and kwargs, changing class names and so on.

I have some more reasons, but maybe the above ones are convincing
enough.

## Minor changes: improved input validation and types

I added input validation for things like `state` and `action_scaling`
(which only makes sense for continuous envs). After adding this, some
tests failed to pass this validation. There I added
`action_scaling=isinstance(env.action_space, Box)`, after which tests
were green. I don't know why the tests were green before, since action
scaling doesn't make sense for discrete actions. I guess some aspect was
not tested and didn't crash.

I also added Literal in some places, in particular for
`action_bound_method`. Now it is no longer allowed to pass an empty
string, instead one should pass `None`. Also here there is input
validation with clear error messages.

@Trinkle23897 The functional tests are green. I didn't want to fix the
formatting, since it will change in the next PR that will solve #914
anyway. I also found a whole bunch of code in `docs/_static`, which I
just deleted (shouldn't it be copied from the sources during docs build
instead of committed?). I also haven't adjusted the documentation yet,
which atm still mentions the trainers of the type
`onpolicy_trainer(...)` instead of `OnpolicyTrainer(...).run()`

## Breaking Changes

The adjustments to the trainer package introduce breaking changes as
duplicated interfaces are deleted. However, it should be very easy for
users to adjust to them

---------

Co-authored-by: Michael Panchenko <m.panchenko@appliedai.de>
2023-08-22 09:54:46 -07:00
..
2023-02-03 11:57:27 -08:00

Offline

In offline reinforcement learning setting, the agent learns a policy from a fixed dataset which is collected once with any policy. And the agent does not interact with environment anymore.

Continuous control

Once the dataset is collected, it will not be changed during training. We use d4rl datasets to train offline agent for continuous control. You can refer to d4rl to see how to use d4rl datasets.

We provide implementation of BCQ and CQL algorithm for continuous control.

Train

Tianshou provides an offline_trainer for offline reinforcement learning. You can parse d4rl datasets into a ReplayBuffer , and set it as the parameter buffer of offline_trainer. d4rl_bcq.py is an example of offline RL using the d4rl dataset.

Results

IL (Imitation Learning, aka, Behavior Cloning)

Environment Dataset IL Parameters
HalfCheetah-v2 halfcheetah-expert-v2 11355.31 python3 d4rl_il.py --task HalfCheetah-v2 --expert-data-task halfcheetah-expert-v2
HalfCheetah-v2 halfcheetah-medium-v2 5098.16 python3 d4rl_il.py --task HalfCheetah-v2 --expert-data-task halfcheetah-medium-v2

BCQ

Environment Dataset BCQ Parameters
HalfCheetah-v2 halfcheetah-expert-v2 11509.95 python3 d4rl_bcq.py --task HalfCheetah-v2 --expert-data-task halfcheetah-expert-v2
HalfCheetah-v2 halfcheetah-medium-v2 5147.43 python3 d4rl_bcq.py --task HalfCheetah-v2 --expert-data-task halfcheetah-medium-v2

CQL

Environment Dataset CQL Parameters
HalfCheetah-v2 halfcheetah-expert-v2 2864.37 python3 d4rl_cql.py --task HalfCheetah-v2 --expert-data-task halfcheetah-expert-v2
HalfCheetah-v2 halfcheetah-medium-v2 6505.41 python3 d4rl_cql.py --task HalfCheetah-v2 --expert-data-task halfcheetah-medium-v2

TD3+BC

Environment Dataset CQL Parameters
HalfCheetah-v2 halfcheetah-expert-v2 11788.25 python3 d4rl_td3_bc.py --task HalfCheetah-v2 --expert-data-task halfcheetah-expert-v2
HalfCheetah-v2 halfcheetah-medium-v2 5741.13 python3 d4rl_td3_bc.py --task HalfCheetah-v2 --expert-data-task halfcheetah-medium-v2

Observation normalization

Following the original paper, we use observation normalization by default. You can turn it off by setting --norm-obs 0. The difference are small but consistent.

Dataset w/ norm-obs w/o norm-obs
halfcheeta-medium-v2 5741.13 5724.41
halfcheeta-expert-v2 11788.25 11665.77
walker2d-medium-v2 4051.76 3985.59
walker2d-expert-v2 5068.15 5027.75

Discrete control

For discrete control, we currently use ad hoc Atari data generated from a trained QRDQN agent.

Gather Data

To running CQL algorithm on Atari, you need to do the following things:

  • Train an expert, by using the command listed in the QRDQN section of Atari examples: python3 atari_qrdqn.py --task {your_task}
  • Generate buffer with noise: python3 atari_qrdqn.py --task {your_task} --watch --resume-path log/{your_task}/qrdqn/policy.pth --eps-test 0.2 --buffer-size 1000000 --save-buffer-name expert.hdf5 (note that 1M Atari buffer cannot be saved as .pkl format because it is too large and will cause error);
  • Train offline model: python3 atari_{bcq,cql,crr}.py --task {your_task} --load-buffer-name expert.hdf5.

IL

We test our IL implementation on two example tasks (different from author's version, we use v4 instead of v0; one epoch means 10k gradient step):

Task Online QRDQN Behavioral IL parameters
PongNoFrameskip-v4 20.5 6.8 20.0 (epoch 5) python3 atari_il.py --task PongNoFrameskip-v4 --load-buffer-name log/PongNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 5
BreakoutNoFrameskip-v4 394.3 46.9 121.9 (epoch 12, could be higher) python3 atari_il.py --task BreakoutNoFrameskip-v4 --load-buffer-name log/BreakoutNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 12

BCQ

We test our BCQ implementation on two example tasks (different from author's version, we use v4 instead of v0; one epoch means 10k gradient step):

Task Online QRDQN Behavioral BCQ parameters
PongNoFrameskip-v4 20.5 6.8 20.1 (epoch 5) python3 atari_bcq.py --task PongNoFrameskip-v4 --load-buffer-name log/PongNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 5
BreakoutNoFrameskip-v4 394.3 46.9 64.6 (epoch 12, could be higher) python3 atari_bcq.py --task BreakoutNoFrameskip-v4 --load-buffer-name log/BreakoutNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 12

CQL

We test our CQL implementation on two example tasks (different from author's version, we use v4 instead of v0; one epoch means 10k gradient step):

Task Online QRDQN Behavioral CQL parameters
PongNoFrameskip-v4 20.5 6.8 20.4 (epoch 5) python3 atari_cql.py --task PongNoFrameskip-v4 --load-buffer-name log/PongNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 5
BreakoutNoFrameskip-v4 394.3 46.9 129.4 (epoch 12) python3 atari_cql.py --task BreakoutNoFrameskip-v4 --load-buffer-name log/BreakoutNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 12 --min-q-weight 50

We reduce the size of the offline data to 10% and 1% of the above and get:

Buffer size 100000:

Task Online QRDQN Behavioral CQL parameters
PongNoFrameskip-v4 20.5 5.8 21 (epoch 5) python3 atari_cql.py --task PongNoFrameskip-v4 --load-buffer-name log/PongNoFrameskip-v4/qrdqn/expert.size_1e5.hdf5 --epoch 5
BreakoutNoFrameskip-v4 394.3 41.4 40.8 (epoch 12) python3 atari_cql.py --task BreakoutNoFrameskip-v4 --load-buffer-name log/BreakoutNoFrameskip-v4/qrdqn/expert.size_1e5.hdf5 --epoch 12 --min-q-weight 20

Buffer size 10000:

Task Online QRDQN Behavioral CQL parameters
PongNoFrameskip-v4 20.5 nan 1.8 (epoch 5) python3 atari_cql.py --task PongNoFrameskip-v4 --load-buffer-name log/PongNoFrameskip-v4/qrdqn/expert.size_1e4.hdf5 --epoch 5 --min-q-weight 1
BreakoutNoFrameskip-v4 394.3 31.7 22.5 (epoch 12) python3 atari_cql.py --task BreakoutNoFrameskip-v4 --load-buffer-name log/BreakoutNoFrameskip-v4/qrdqn/expert.size_1e4.hdf5 --epoch 12 --min-q-weight 10

CRR

We test our CRR implementation on two example tasks (different from author's version, we use v4 instead of v0; one epoch means 10k gradient step):

Task Online QRDQN Behavioral CRR CRR w/ CQL parameters
PongNoFrameskip-v4 20.5 6.8 -21 (epoch 5) 17.7 (epoch 5) python3 atari_crr.py --task PongNoFrameskip-v4 --load-buffer-name log/PongNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 5
BreakoutNoFrameskip-v4 394.3 46.9 23.3 (epoch 12) 76.9 (epoch 12) python3 atari_crr.py --task BreakoutNoFrameskip-v4 --load-buffer-name log/BreakoutNoFrameskip-v4/qrdqn/expert.hdf5 --epoch 12 --min-q-weight 50

Note that CRR itself does not work well in Atari tasks but adding CQL loss/regularizer helps.

RL Unplugged Data

We provide a script to convert the Atari datasets of RL Unplugged to Tianshou ReplayBuffer.

For example, the following command will download the first shard of the first run of Breakout game to ~/.rl_unplugged/datasets/Breakout/run_1-00001-of-00100 then convert it to a tianshou.data.ReplayBuffer and save it to ~/.rl_unplugged/buffers/Breakout/run_1-00001-of-00100.hdf5 (use --dataset-dir and --buffer-dir to change the default directories):

python3 convert_rl_unplugged_atari.py --task Breakout --run-id 1 --shard-id 1

Then you can use it to train an agent by:

python3 atari_bcq.py --task BreakoutNoFrameskip-v4 --load-buffer-name ~/.rl_unplugged/datasets/Breakout/run_1-00001-of-00100.hdf5 --buffer-from-rl-unplugged --epoch 12

Note:

  • Each shard contains about 500k transitions.
  • This conversion script depends on Tensorflow.
  • It takes about 1 hour to process one shard on my machine. YMMV.