Update

[ghstack-poisoned]

Author: vmoens

torchrl/data/datasets/d4rl.py

@@ -279,6 +279,7 @@ def _get_dataset_direct(self, name, env_kwargs):
         # so we need to ensure we're using the gym backend
         with set_gym_backend("gym"):
             import gym
+
             env = GymWrapper(gym.make(name))
         with tempfile.TemporaryDirectory() as tmpdir:
             os.environ["D4RL_DATASET_DIR"] = tmpdir
@@ -358,6 +359,7 @@ def _get_dataset_from_env(self, name, env_kwargs):
         # so we need to ensure we're using the gym backend
         with set_gym_backend("gym"), tempfile.TemporaryDirectory() as tmpdir:
             import gym
+
             os.environ["D4RL_DATASET_DIR"] = tmpdir
             env = GymWrapper(gym.make(name))
             dataset = make_tensordict(

torchrl/envs/transforms/module.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
 
 from collections.abc import Callable
 from contextlib import nullcontext