Update

[ghstack-poisoned]

Author: Vincent Moens

setup.cfg

@@ -45,3 +45,7 @@ ignore-decorators =
     test_*
 ;    test/*.py
 ;    .circleci/*
+
+[autoflake]
+per-file-ignores =
+    torchrl/trainers/helpers/envs.py *

test/test_helpers.py

@@ -2,6 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
 
 import argparse
 import dataclasses
@@ -225,6 +226,7 @@ def test_timeit():
 @pytest.mark.skipif(not _has_hydra, reason="No hydra library found")
 @pytest.mark.parametrize("from_pixels", [(), ("from_pixels=True", "catframes=4")])
 def test_transformed_env_constructor_with_state_dict(from_pixels):
+
     config_fields = [
         (config_field.name, config_field.type, config_field)
         for config_cls in (

torchrl/objectives/common.py

@@ -21,7 +21,7 @@
 
 from torchrl._utils import RL_WARNINGS
 from torchrl.envs.utils import ExplorationType, set_exploration_type
-from torchrl.modules import set_recurrent_mode
+from torchrl.modules.tensordict_module.rnn import set_recurrent_mode
 from torchrl.objectives.utils import ValueEstimators
 from torchrl.objectives.value import ValueEstimatorBase
 

torchrl/trainers/helpers/envs.py

@@ -2,13 +2,17 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-from __future__ import annotations
+
+# This makes omegaconf unhappy with typing.Any
+# Therefore we need Optional and Union
+# from __future__ import annotations
 
 from copy import copy
 from dataclasses import dataclass, field as dataclass_field
-from typing import Any, Callable, Sequence
+from typing import Any, Callable, Optional, Sequence, Union
 
 import torch
+from omegaconf import DictConfig
 
 from torchrl._utils import logger as torchrl_logger, VERBOSE
 from torchrl.envs import ParallelEnv
@@ -212,18 +216,18 @@ def get_norm_state_dict(env):
 def transformed_env_constructor(
     cfg: DictConfig,  # noqa: F821
     video_tag: str = "",
-    logger: Logger | None = None,
-    stats: dict | None = None,
+    logger: Optional[Logger] = None,  # noqa
+    stats: Optional[dict] = None,
     norm_obs_only: bool = False,
     use_env_creator: bool = False,
-    custom_env_maker: Callable | None = None,
-    custom_env: EnvBase | None = None,
+    custom_env_maker: Optional[Callable] = None,
+    custom_env: Optional[EnvBase] = None,
     return_transformed_envs: bool = True,
-    action_dim_gsde: int | None = None,
-    state_dim_gsde: int | None = None,
-    batch_dims: int | None = 0,
-    obs_norm_state_dict: dict | None = None,
-) -> Callable | EnvCreator:
+    action_dim_gsde: Optional[int] = None,
+    state_dim_gsde: Optional[int] = None,
+    batch_dims: Optional[int] = 0,
+    obs_norm_state_dict: Optional[dict] = None,
+) -> Union[Callable, EnvCreator]:
     """Returns an environment creator from an argparse.Namespace built with the appropriate parser constructor.
 
     Args:
@@ -329,7 +333,7 @@ def make_transformed_env(**kwargs) -> TransformedEnv:
 
 def parallel_env_constructor(
     cfg: DictConfig, **kwargs  # noqa: F821
-) -> ParallelEnv | EnvCreator:
+) -> Union[ParallelEnv, EnvCreator]:
     """Returns a parallel environment from an argparse.Namespace built with the appropriate parser constructor.
 
     Args:
@@ -374,7 +378,7 @@ def parallel_env_constructor(
 def get_stats_random_rollout(
     cfg: DictConfig,  # noqa: F821
     proof_environment: EnvBase = None,
-    key: str | None = None,
+    key: Optional[str] = None,
 ):
     """Gathers stas (loc and scale) from an environment using random rollouts.
 
@@ -452,7 +456,7 @@ def get_stats_random_rollout(
 def initialize_observation_norm_transforms(
     proof_environment: EnvBase,
     num_iter: int = 1000,
-    key: str | tuple[str, ...] = None,
+    key: Union[str, tuple[str, ...]] = None,
 ):
     """Calls :obj:`ObservationNorm.init_stats` on all uninitialized :obj:`ObservationNorm` instances of a :obj:`TransformedEnv`.
 
@@ -532,7 +536,7 @@ class EnvConfig:
     # maximum steps per trajectory, frames per batch or any other factor in the algorithm,
     # e.g. if the total number of frames that has to be computed is 50e6 and the frame skip is 4
     # the actual number of frames retrieved will be 200e6. Default=1.
-    reward_scaling: float | None = None
+    reward_scaling: Any = None  # noqa
     # scale of the reward.
     reward_loc: float = 0.0
     # location of the reward.