Fix OpenEnv CI collection and formatting

Author: vmoens

test/libs/test_openenv.py

@@ -12,12 +12,11 @@
 
 import pytest
 import torch
-from omegaconf import DictConfig
-from tensordict import TensorDict
-from tensordict.tensorclass import NonTensorData
 
 import torchrl.envs.libs.openenv as openenv_mod
 import torchrl.envs.llm.libs.openenv as openenv_chat_mod
+from tensordict import TensorDict
+from tensordict.tensorclass import NonTensorData
 from torchrl.data import LazyStackStorage, ReplayBuffer
 from torchrl.data.llm import History
 from torchrl.envs.libs.openenv import OpenEnvEnv, OpenEnvWrapper
@@ -26,6 +25,8 @@
 from torchrl.modules.llm.policies.common import ChatHistory
 from torchrl.objectives.llm.grpo import MCAdvantage
 
+_has_omegaconf = importlib.util.find_spec("omegaconf") is not None
+
 
 @dataclass
 class _StepResult:
@@ -46,6 +47,23 @@ def model_dump(self):
         return {"prompt": ["nested", {"value": 1}], "reward": 2.0, "done": True}
 
 
+class _Config(dict):
+    def __init__(self, data):
+        super().__init__((key, self._convert(value)) for key, value in data.items())
+
+    @classmethod
+    def _convert(cls, value):
+        if isinstance(value, dict):
+            return cls(value)
+        return value
+
+    def __getattr__(self, name):
+        try:
+            return self[name]
+        except KeyError as err:
+            raise AttributeError(name) from err
+
+
 class _SyncOpenEnv:
     def __init__(self):
         self.connected = False
@@ -226,6 +244,8 @@ def test_rand_step_check_env_specs_and_rollout(self):
 
 class TestOpenEnvGRPO:
     def test_make_env_openenv_with_local_fixture(self, monkeypatch):
+        if not _has_omegaconf:
+            pytest.skip("omegaconf is required to import the GRPO recipe helpers")
         pytest.importorskip("transformers")
         pytest.importorskip("openenv")
         spec = importlib.util.spec_from_file_location(
@@ -246,7 +266,7 @@ def test_make_env_openenv_with_local_fixture(self, monkeypatch):
             "from_env",
             staticmethod(lambda name: _TextAction),
         )
-        cfg = DictConfig(
+        cfg = _Config(
             {
                 "env": {
                     "dataset": "openenv",

torchrl/data/llm/__init__.py

@@ -3,13 +3,13 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .history import add_chat_template, ContentBase, History
 from .dataset import (
     create_infinite_iterator,
     get_dataloader,
     TensorDictTokenizer,
     TokenizedDatasetLoader,
 )
+from .history import add_chat_template, ContentBase, History
 from .prompt import PromptData, PromptTensorDictTokenizer
 from .reward import PairwiseDataset, RewardData
 from .topk import TopKRewardSelector

torchrl/envs/llm/chat.py

@@ -12,15 +12,14 @@
 from tensordict import lazy_stack, TensorDictBase
 from tensordict.utils import _zip_strict
 from torch.utils.data import DataLoader
-from torchrl.data.tensor_specs import Composite, NonTensor
 from torchrl.data.llm.history import History
-from torchrl.envs.common import EnvBase
-from torchrl.envs.transforms import TransformedEnv
-from torchrl.envs.common import _EnvPostInit
+from torchrl.data.tensor_specs import Composite, NonTensor
+from torchrl.envs.common import _EnvPostInit, EnvBase
 from torchrl.envs.llm.transforms.dataloading import (
     DataLoadingPrimer,
     RayDataLoadingPrimer,
 )
+from torchrl.envs.transforms import TransformedEnv
 from torchrl.modules.llm.policies.common import ChatHistory, Text, Tokens
 
 if TYPE_CHECKING:

torchrl/envs/llm/datasets/gsm8k.py

@@ -14,12 +14,12 @@
 from tensordict.utils import _zip_strict
 from torch.utils.data import DataLoader
 from torchrl.data import TensorSpec
-from torchrl.envs.transforms import StepCounter, Transform
 
 from torchrl.envs.llm.chat import DatasetChatEnv
 
 from torchrl.envs.llm.envs import LLMEnv
 from torchrl.envs.llm.reward.gsm8k import GSM8KRewardParser
+from torchrl.envs.transforms import StepCounter, Transform
 
 if TYPE_CHECKING:
     import transformers

torchrl/envs/llm/datasets/ifeval.py

@@ -11,9 +11,9 @@
 import torch
 from tensordict import NonTensorData, NonTensorStack, TensorClass, TensorDict
 from torchrl.data import Composite, NonTensor, Unbounded
-from torchrl.envs.transforms import StepCounter
 from torchrl.envs.llm.chat import DatasetChatEnv
 from torchrl.envs.llm.reward.ifeval import IfEvalScorer
+from torchrl.envs.transforms import StepCounter
 
 if TYPE_CHECKING:
     import transformers

torchrl/envs/llm/libs/openenv.py

@@ -276,7 +276,9 @@ def _format_observation(self, observation: Any) -> Any:
         )
 
     def _wrap_observation(self, observation: Any) -> NonTensorData:
-        return NonTensorData(observation, batch_size=self.batch_size, device=self.device)
+        return NonTensorData(
+            observation, batch_size=self.batch_size, device=self.device
+        )
 
     def _make_history_message(self, role: str, content: Any) -> History:
         return History(

torchrl/envs/llm/reward/gsm8k.py

@@ -11,8 +11,8 @@
 from tensordict import lazy_stack, NestedKey, TensorDict, TensorDictBase
 from tensordict.utils import _zip_strict, is_non_tensor
 from torchrl.data import Composite, Unbounded
-from torchrl.envs.transforms import Transform
 from torchrl.envs.common import EnvBase
+from torchrl.envs.transforms import Transform
 
 
 class GSM8KRewardParser(Transform):

torchrl/envs/llm/transforms/reason.py

@@ -13,8 +13,8 @@
 from torchrl._utils import logger as torchrl_logger
 
 from torchrl.data.llm.history import History
-from torchrl.envs.transforms import Transform
 from torchrl.envs.common import EnvBase
+from torchrl.envs.transforms import Transform
 
 
 class AddThinkingPrompt(Transform):