Add Type Annotations

[frederikschubert]

This issue tracks the progress of adding type annotations to coax.

• [ ] _core
• [ ] experience_replay
• [ ] model_updaters
• [ ] policy_objectives
• [ ] proba_dists
• [ ] reward_tracing
• [ ] td_learning
• [ ] utils
• [ ] value_transforms
• [ ] wrappers

The types are added by utilising pyannotate and adding the following snippet to the coax._base.TestCase class:

...
@classmethod
    def setUpClass(cls) -> None:
        collect_types.init_types_collection()
        collect_types.start()

    @classmethod
    def tearDownClass(cls) -> None:
        collect_types.stop()
        type_replacements = {
            "jaxlib.xla_extension.DeviceArray": "jax.numpy.ndarray",
            "haiku._src.data_structures.FlatMapping": "typing.Mapping",
            "coax._core.policy_test": "gym.Env"
        }
        types_str = collect_types.dumps_stats()
        for inferred_type, replacement in type_replacements.items():
            types_str = types_str.replace(inferred_type, replacement)
        with open(sys.modules[cls.__module__].__file__.replace(".py", "_types.json"), "w") as f:
            f.write(types_str)
...

and the types are added automatically

for t in coax/**/*_test_types.json
do
    pyannotate --type-info $t -3 coax/* -w
done