We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] importing `dm_control` environments when the library is not installed does not display the right error message
## Describe the bug

We allow users to load the file `envs/libs/dm_control.py` even if `dm-control` is not installes. When creating an object in this situation, the error message now is that the environment configuration cannot be found. An `ImportError` should be raised instead (after a check for `_has_dmc` has been made).

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchrl/envs/libs/dm_control.py`
Content:
```
1 # Copyright (c) Meta Platforms, Inc. and affiliates.
2 #
3 # This source code is licensed under the MIT license found in the
4 # LICENSE file in the root directory of this source tree.
5 from __future__ import annotations
6 
7 import collections
8 import os
9 from typing import Optional, Tuple, Union, Dict, Any
10 
11 import numpy as np
12 import torch
13 
14 from torchrl.data import (
15     CompositeSpec,
16     NdBoundedTensorSpec,
17     NdUnboundedContinuousTensorSpec,
18     TensorSpec,
19     NdUnboundedDiscreteTensorSpec,
20 )
21 from ...data.utils import numpy_to_torch_dtype_dict, DEVICE_TYPING
22 from ..gym_like import GymLikeEnv
23 
24 if torch.has_cuda and torch.cuda.device_count() > 1:
25     n = torch.cuda.device_count() - 1
26     os.environ["EGL_DEVICE_ID"] = str(1 + (os.getpid() % n))
27     print("EGL_DEVICE_ID: ", os.environ["EGL_DEVICE_ID"])
28 
29 try:
30 
31     import dm_control
32     import dm_env
33     from dm_control import suite
34     from dm_control.suite.wrappers import pixels
35 
36     _has_dmc = True
37 
38 except ImportError:
39     _has_dmc = False
40 
41 __all__ = ["DMControlEnv", "DMControlWrapper"]
42 
43 
44 def _dmcontrol_to_torchrl_spec_transform(
45     spec,
46     dtype: Optional[torch.dtype] = None,
47     device: DEVICE_TYPING = None,
48 ) -> TensorSpec:
49     if isinstance(spec, collections.OrderedDict):
50         spec = {
51             "next_" + k: _dmcontrol_to_torchrl_spec_transform(item, device=device)
52             for k, item in spec.items()
53         }
54         return CompositeSpec(**spec)
55     elif isinstance(spec, dm_env.specs.BoundedArray):
56         if dtype is None:
57             dtype = numpy_to_torch_dtype_dict[spec.dtype]
58         return NdBoundedTensorSpec(
59             shape=spec.shape,
60             minimum=spec.minimum,
61             maximum=spec.maximum,
62             dtype=dtype,
63             device=device,
64         )
65     elif isinstance(spec, dm_env.specs.Array):
66         if dtype is None:
67             dtype = numpy_to_torch_dtype_dict[spec.dtype]
68         if dtype in (torch.float, torch.double, torch.half):
69             return NdUnboundedContinuousTensorSpec(
70                 shape=spec.shape, dtype=dtype, device=device
71             )
72         else:
73             return NdUnboundedDiscreteTensorSpec(
74                 shape=spec.shape, dtype=dtype, device=device
75             )
76 
77     else:
78         raise NotImplementedError(type(spec))
79 
80 
81 def _get_envs(to_dict: bool = True) -> Dict[str, Any]:
82     if not _has_dmc:
83         return dict()
84     if not to_dict:
85         return tuple(suite.BENCHMARKING) + tuple(suite.EXTRA)
86     d = dict()
87     for tup in suite.BENCHMARKING:
88         env_name = tup[0]
89         d.setdefault(env_name, []).append(tup[1])
90     for tup in suite.EXTRA:
91         env_name = tup[0]
92         d.setdefault(env_name, []).append(tup[1])
93     return d
94 
95 
96 def _robust_to_tensor(array: Union[float, np.ndarray]) -> torch.Tensor:
97     if isinstance(array, np.ndarray):
98         return torch.tensor(array.copy())
99     else:
100         return torch.tensor(array)
101 
102 
103 class DMControlWrapper(GymLikeEnv):
104     """
105     DeepMind Control lab environment wrapper.
106 
107     Args:
108         env (dm_control.suite env): environment instance
109         from_pixels (bool): if True, the observation
110 
111     Examples:
112         >>> env = dm_control.suite.load("cheetah", "run")
113         >>> env = DMControlWrapper(env,
114         ...    from_pixels=True, frame_skip=4)
115         >>> td = env.rand_step()
116         >>> print(td)
117         >>> print(env.available_envs)
118     """
119 
120     git_url = "https://github.com/deepmind/dm_control"
121     libname = "dm_control"
122     available_envs = _get_envs()
123 
124     def __init__(self, env=None, **kwargs):
125         if env is not None:
126             kwargs["env"] = env
127         super().__init__(**kwargs)
128 
129     def _build_env(
130         self,
131         env,
132         _seed: Optional[int] = None,
133         from_pixels: bool = False,
134         render_kwargs: Optional[dict] = None,
135         pixels_only: bool = False,
136         **kwargs,
137     ):
138         self.from_pixels = from_pixels
139         self.pixels_only = pixels_only
140 
141         if from_pixels:
142             self._set_egl_device(self.device)
143             self.render_kwargs = {"camera_id": 0}
144             if render_kwargs is not None:
145                 self.render_kwargs.update(render_kwargs)
146             env = pixels.Wrapper(
147                 env,
148                 pixels_only=self.pixels_only,
149                 render_kwargs=self.render_kwargs,
150             )
151         return env
152 
153     def _make_specs(self, env: "gym.Env") -> None:  # noqa: F821
154         # specs are defined when first called
155         return
156 
157     def _check_kwargs(self, kwargs: Dict):
158         if "env" not in kwargs:
159             raise TypeError("Could not find environment key 'env' in kwargs.")
160         env = kwargs["env"]
161         if not isinstance(env, (dm_control.rl.control.Environment, pixels.Wrapper)):
162             raise TypeError(
163                 "env is not of type 'dm_control.rl.control.Environment' or `dm_control.suite.wrappers.pixels.Wrapper`."
164             )
165 
166     def _set_egl_device(self, device: DEVICE_TYPING):
167         # Deprecated as lead to unreliable rendering
168         # egl device needs to be set before importing mujoco bindings: in
169         # distributed settings, it'll be easy to tell which cuda device to use.
170         # In mp settings, we'll need to use mp.Pool with a specific init function
171         # that defines the EGL device before importing libraries. For now, we'll
172         # just use a common EGL_DEVICE_ID environment variable for all processes.
173         return
174 
175     def to(self, device: DEVICE_TYPING) -> DMControlEnv:
176         super().to(device)
177         self._set_egl_device(self.device)
178         return self
179 
180     def _init_env(self, seed: Optional[int] = None) -> Optional[int]:
181         seed = self.set_seed(seed)
182         return seed
183 
184     def _set_seed(self, _seed: Optional[int]) -> Optional[int]:
185         if _seed is None:
186             return None
187         random_state = np.random.RandomState(_seed)
188         if isinstance(self._env, pixels.Wrapper):
189             if not hasattr(self._env._env.task, "_random"):
190                 raise RuntimeError("self._env._env.task._random does not exist")
191             self._env._env.task._random = random_state
192         else:
193             if not hasattr(self._env.task, "_random"):
194                 raise RuntimeError("self._env._env.task._random does not exist")
195             self._env.task._random = random_state
196         self.reset()
197         return _seed
198 
199     def _output_transform(
200         self, timestep_tuple: Tuple["TimeStep"]  # noqa: F821
201     ) -> Tuple[np.ndarray, float, bool]:
202         if type(timestep_tuple) is not tuple:
203             timestep_tuple = (timestep_tuple,)
204         reward = timestep_tuple[0].reward
205 
206         done = False  # dm_control envs are non-terminating
207         observation = timestep_tuple[0].observation
208         return observation, reward, done
209 
210     @property
211     def input_spec(self) -> TensorSpec:
212         if self._input_spec is None:
213             self._input_spec = CompositeSpec(
214                 action=_dmcontrol_to_torchrl_spec_transform(
215                     self._env.action_spec(), device=self.device
216                 )
217             )
218         return self._input_spec
219 
220     @input_spec.setter
221     def input_spec(self, value: TensorSpec) -> None:
222         self._input_spec = value
223 
224     @property
225     def observation_spec(self) -> TensorSpec:
226         if self._observation_spec is None:
227             self._observation_spec = _dmcontrol_to_torchrl_spec_transform(
228                 self._env.observation_spec(), device=self.device
229             )
230         return self._observation_spec
231 
232     @observation_spec.setter
233     def observation_spec(self, value: TensorSpec) -> None:
234         self._observation_spec = value
235 
236     @property
237     def reward_spec(self) -> TensorSpec:
238         if self._reward_spec is None:
239             self._reward_spec = _dmcontrol_to_torchrl_spec_transform(
240                 self._env.reward_spec(), device=self.device
241             )
242         return self._reward_spec
243 
244     @reward_spec.setter
245     def reward_spec(self, value: TensorSpec) -> None:
246         self._reward_spec = value
247 
248     def __repr__(self) -> str:
249         return (
250             f"{self.__class__.__name__}(env={self._env}, batch_size={self.batch_size})"
251         )
252 
253 
254 class DMControlEnv(DMControlWrapper):
255     """
256     DeepMind Control lab environment wrapper.
257 
258     Args:
259         env_name (str): name of the environment
260         task_name (str): name of the task
261         seed (int, optional): seed to use for the environment
262         from_pixels (bool, optional): if True, the observation will be returned
263             as an image.
264             Default is False.
265 
266     Examples:
267         >>> env = DMControlEnv(env_name="cheetah", task_name="run",
268         ...    from_pixels=True, frame_skip=4)
269         >>> td = env.rand_step()
270         >>> print(td)
271         >>> print(env.available_envs)
272     """
273 
274     def __init__(self, env_name, task_name, **kwargs):
275         kwargs["env_name"] = env_name
276         kwargs["task_name"] = task_name
277         super().__init__(**kwargs)
278 
279     def _build_env(
280         self,
281         env_name: str,
282         task_name: str,
283         _seed: Optional[int] = None,
284         **kwargs,
285     ):
286         self.env_name = env_name
287         self.task_name = task_name
288 
289         from_pixels = kwargs.get("from_pixels")
290         if "from_pixels" in kwargs:
291             del kwargs["from_pixels"]
292         pixels_only = kwargs.get("pixels_only")
293         if "pixels_only" in kwargs:
294             del kwargs["pixels_only"]
295 
296         if not _has_dmc:
297             raise RuntimeError(
298                 f"dm_control not found, unable to create {env_name}:"
299                 f" {task_name}. Consider downloading and installing "
300                 f"dm_control from {self.git_url}"
301             )
302 
303         if _seed is not None:
304             random_state = np.random.RandomState(_seed)
305             kwargs = {"random": random_state}
306         env = suite.load(env_name, task_name, task_kwargs=kwargs)
307         return super()._build_env(
308             env, from_pixels=from_pixels, pixels_only=pixels_only, **kwargs
309         )
310 
311     def rebuild_with_kwargs(self, **new_kwargs):
312         self._constructor_kwargs.update(new_kwargs)
313         self._env = self._build_env()
314         self._make_specs(self._env)
315 
316     def _check_kwargs(self, kwargs: Dict):
317         if "env_name" in kwargs:
318             env_name = kwargs["env_name"]
319             if "task_name" in kwargs:
320                 task_name = kwargs["task_name"]
321                 if (
322                     env_name not in self.available_envs
323                     or task_name not in self.available_envs[env_name]
324                 ):
325                     raise RuntimeError(
326                         f"{env_name} with task {task_name} is unknown in {self.libname}"
327                     )
328             else:
329                 raise TypeError("dm_control requires task_name to be specified")
330         else:
331             raise TypeError("dm_control requires env_name to be specified")
332 
333     # def _set_seed(self, _seed: int) -> int:
334     #     self._env = self._build_env(
335     #         _seed=_seed, **self._constructor_kwargs
336     #     )
337     #     self.reset()
338     #     return _seed
339 
340     def __repr__(self) -> str:
341         return f"{self.__class__.__name__}(env={self.env_name}, task={self.task_name}, batch_size={self.batch_size})"
342
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torchrl/envs/libs/dm_control.py b/torchrl/envs/libs/dm_control.py
--- a/torchrl/envs/libs/dm_control.py
+++ b/torchrl/envs/libs/dm_control.py
@@ -272,6 +272,11 @@
     """
 
     def __init__(self, env_name, task_name, **kwargs):
+        if not _has_dmc:
+            raise ImportError(
+                "dm_control python package was not found."
+                "Please install this dependency."
+            )
         kwargs["env_name"] = env_name
         kwargs["task_name"] = task_name
         super().__init__(**kwargs)

{"golden_diff": "diff --git a/torchrl/envs/libs/dm_control.py b/torchrl/envs/libs/dm_control.py\n--- a/torchrl/envs/libs/dm_control.py\n+++ b/torchrl/envs/libs/dm_control.py\n@@ -272,6 +272,11 @@\n     \"\"\"\n \n     def __init__(self, env_name, task_name, **kwargs):\n+        if not _has_dmc:\n+            raise ImportError(\n+                \"dm_control python package was not found.\"\n+                \"Please install this dependency.\"\n+            )\n         kwargs[\"env_name\"] = env_name\n         kwargs[\"task_name\"] = task_name\n         super().__init__(**kwargs)\n", "issue": "[BUG] importing `dm_control` environments when the library is not installed does not display the right error message\n## Describe the bug\r\n\r\nWe allow users to load the file `envs/libs/dm_control.py` even if `dm-control` is not installes. When creating an object in this situation, the error message now is that the environment configuration cannot be found. An `ImportError` should be raised instead (after a check for `_has_dmc` has been made).\n", "before_files": [{"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\nfrom __future__ import annotations\n\nimport collections\nimport os\nfrom typing import Optional, Tuple, Union, Dict, Any\n\nimport numpy as np\nimport torch\n\nfrom torchrl.data import (\n    CompositeSpec,\n    NdBoundedTensorSpec,\n    NdUnboundedContinuousTensorSpec,\n    TensorSpec,\n    NdUnboundedDiscreteTensorSpec,\n)\nfrom ...data.utils import numpy_to_torch_dtype_dict, DEVICE_TYPING\nfrom ..gym_like import GymLikeEnv\n\nif torch.has_cuda and torch.cuda.device_count() > 1:\n    n = torch.cuda.device_count() - 1\n    os.environ[\"EGL_DEVICE_ID\"] = str(1 + (os.getpid() % n))\n    print(\"EGL_DEVICE_ID: \", os.environ[\"EGL_DEVICE_ID\"])\n\ntry:\n\n    import dm_control\n    import dm_env\n    from dm_control import suite\n    from dm_control.suite.wrappers import pixels\n\n    _has_dmc = True\n\nexcept ImportError:\n    _has_dmc = False\n\n__all__ = [\"DMControlEnv\", \"DMControlWrapper\"]\n\n\ndef _dmcontrol_to_torchrl_spec_transform(\n    spec,\n    dtype: Optional[torch.dtype] = None,\n    device: DEVICE_TYPING = None,\n) -> TensorSpec:\n    if isinstance(spec, collections.OrderedDict):\n        spec = {\n            \"next_\" + k: _dmcontrol_to_torchrl_spec_transform(item, device=device)\n            for k, item in spec.items()\n        }\n        return CompositeSpec(**spec)\n    elif isinstance(spec, dm_env.specs.BoundedArray):\n        if dtype is None:\n            dtype = numpy_to_torch_dtype_dict[spec.dtype]\n        return NdBoundedTensorSpec(\n            shape=spec.shape,\n            minimum=spec.minimum,\n            maximum=spec.maximum,\n            dtype=dtype,\n            device=device,\n        )\n    elif isinstance(spec, dm_env.specs.Array):\n        if dtype is None:\n            dtype = numpy_to_torch_dtype_dict[spec.dtype]\n        if dtype in (torch.float, torch.double, torch.half):\n            return NdUnboundedContinuousTensorSpec(\n                shape=spec.shape, dtype=dtype, device=device\n            )\n        else:\n            return NdUnboundedDiscreteTensorSpec(\n                shape=spec.shape, dtype=dtype, device=device\n            )\n\n    else:\n        raise NotImplementedError(type(spec))\n\n\ndef _get_envs(to_dict: bool = True) -> Dict[str, Any]:\n    if not _has_dmc:\n        return dict()\n    if not to_dict:\n        return tuple(suite.BENCHMARKING) + tuple(suite.EXTRA)\n    d = dict()\n    for tup in suite.BENCHMARKING:\n        env_name = tup[0]\n        d.setdefault(env_name, []).append(tup[1])\n    for tup in suite.EXTRA:\n        env_name = tup[0]\n        d.setdefault(env_name, []).append(tup[1])\n    return d\n\n\ndef _robust_to_tensor(array: Union[float, np.ndarray]) -> torch.Tensor:\n    if isinstance(array, np.ndarray):\n        return torch.tensor(array.copy())\n    else:\n        return torch.tensor(array)\n\n\nclass DMControlWrapper(GymLikeEnv):\n    \"\"\"\n    DeepMind Control lab environment wrapper.\n\n    Args:\n        env (dm_control.suite env): environment instance\n        from_pixels (bool): if True, the observation\n\n    Examples:\n        >>> env = dm_control.suite.load(\"cheetah\", \"run\")\n        >>> env = DMControlWrapper(env,\n        ...    from_pixels=True, frame_skip=4)\n        >>> td = env.rand_step()\n        >>> print(td)\n        >>> print(env.available_envs)\n    \"\"\"\n\n    git_url = \"https://github.com/deepmind/dm_control\"\n    libname = \"dm_control\"\n    available_envs = _get_envs()\n\n    def __init__(self, env=None, **kwargs):\n        if env is not None:\n            kwargs[\"env\"] = env\n        super().__init__(**kwargs)\n\n    def _build_env(\n        self,\n        env,\n        _seed: Optional[int] = None,\n        from_pixels: bool = False,\n        render_kwargs: Optional[dict] = None,\n        pixels_only: bool = False,\n        **kwargs,\n    ):\n        self.from_pixels = from_pixels\n        self.pixels_only = pixels_only\n\n        if from_pixels:\n            self._set_egl_device(self.device)\n            self.render_kwargs = {\"camera_id\": 0}\n            if render_kwargs is not None:\n                self.render_kwargs.update(render_kwargs)\n            env = pixels.Wrapper(\n                env,\n                pixels_only=self.pixels_only,\n                render_kwargs=self.render_kwargs,\n            )\n        return env\n\n    def _make_specs(self, env: \"gym.Env\") -> None:  # noqa: F821\n        # specs are defined when first called\n        return\n\n    def _check_kwargs(self, kwargs: Dict):\n        if \"env\" not in kwargs:\n            raise TypeError(\"Could not find environment key 'env' in kwargs.\")\n        env = kwargs[\"env\"]\n        if not isinstance(env, (dm_control.rl.control.Environment, pixels.Wrapper)):\n            raise TypeError(\n                \"env is not of type 'dm_control.rl.control.Environment' or `dm_control.suite.wrappers.pixels.Wrapper`.\"\n            )\n\n    def _set_egl_device(self, device: DEVICE_TYPING):\n        # Deprecated as lead to unreliable rendering\n        # egl device needs to be set before importing mujoco bindings: in\n        # distributed settings, it'll be easy to tell which cuda device to use.\n        # In mp settings, we'll need to use mp.Pool with a specific init function\n        # that defines the EGL device before importing libraries. For now, we'll\n        # just use a common EGL_DEVICE_ID environment variable for all processes.\n        return\n\n    def to(self, device: DEVICE_TYPING) -> DMControlEnv:\n        super().to(device)\n        self._set_egl_device(self.device)\n        return self\n\n    def _init_env(self, seed: Optional[int] = None) -> Optional[int]:\n        seed = self.set_seed(seed)\n        return seed\n\n    def _set_seed(self, _seed: Optional[int]) -> Optional[int]:\n        if _seed is None:\n            return None\n        random_state = np.random.RandomState(_seed)\n        if isinstance(self._env, pixels.Wrapper):\n            if not hasattr(self._env._env.task, \"_random\"):\n                raise RuntimeError(\"self._env._env.task._random does not exist\")\n            self._env._env.task._random = random_state\n        else:\n            if not hasattr(self._env.task, \"_random\"):\n                raise RuntimeError(\"self._env._env.task._random does not exist\")\n            self._env.task._random = random_state\n        self.reset()\n        return _seed\n\n    def _output_transform(\n        self, timestep_tuple: Tuple[\"TimeStep\"]  # noqa: F821\n    ) -> Tuple[np.ndarray, float, bool]:\n        if type(timestep_tuple) is not tuple:\n            timestep_tuple = (timestep_tuple,)\n        reward = timestep_tuple[0].reward\n\n        done = False  # dm_control envs are non-terminating\n        observation = timestep_tuple[0].observation\n        return observation, reward, done\n\n    @property\n    def input_spec(self) -> TensorSpec:\n        if self._input_spec is None:\n            self._input_spec = CompositeSpec(\n                action=_dmcontrol_to_torchrl_spec_transform(\n                    self._env.action_spec(), device=self.device\n                )\n            )\n        return self._input_spec\n\n    @input_spec.setter\n    def input_spec(self, value: TensorSpec) -> None:\n        self._input_spec = value\n\n    @property\n    def observation_spec(self) -> TensorSpec:\n        if self._observation_spec is None:\n            self._observation_spec = _dmcontrol_to_torchrl_spec_transform(\n                self._env.observation_spec(), device=self.device\n            )\n        return self._observation_spec\n\n    @observation_spec.setter\n    def observation_spec(self, value: TensorSpec) -> None:\n        self._observation_spec = value\n\n    @property\n    def reward_spec(self) -> TensorSpec:\n        if self._reward_spec is None:\n            self._reward_spec = _dmcontrol_to_torchrl_spec_transform(\n                self._env.reward_spec(), device=self.device\n            )\n        return self._reward_spec\n\n    @reward_spec.setter\n    def reward_spec(self, value: TensorSpec) -> None:\n        self._reward_spec = value\n\n    def __repr__(self) -> str:\n        return (\n            f\"{self.__class__.__name__}(env={self._env}, batch_size={self.batch_size})\"\n        )\n\n\nclass DMControlEnv(DMControlWrapper):\n    \"\"\"\n    DeepMind Control lab environment wrapper.\n\n    Args:\n        env_name (str): name of the environment\n        task_name (str): name of the task\n        seed (int, optional): seed to use for the environment\n        from_pixels (bool, optional): if True, the observation will be returned\n            as an image.\n            Default is False.\n\n    Examples:\n        >>> env = DMControlEnv(env_name=\"cheetah\", task_name=\"run\",\n        ...    from_pixels=True, frame_skip=4)\n        >>> td = env.rand_step()\n        >>> print(td)\n        >>> print(env.available_envs)\n    \"\"\"\n\n    def __init__(self, env_name, task_name, **kwargs):\n        kwargs[\"env_name\"] = env_name\n        kwargs[\"task_name\"] = task_name\n        super().__init__(**kwargs)\n\n    def _build_env(\n        self,\n        env_name: str,\n        task_name: str,\n        _seed: Optional[int] = None,\n        **kwargs,\n    ):\n        self.env_name = env_name\n        self.task_name = task_name\n\n        from_pixels = kwargs.get(\"from_pixels\")\n        if \"from_pixels\" in kwargs:\n            del kwargs[\"from_pixels\"]\n        pixels_only = kwargs.get(\"pixels_only\")\n        if \"pixels_only\" in kwargs:\n            del kwargs[\"pixels_only\"]\n\n        if not _has_dmc:\n            raise RuntimeError(\n                f\"dm_control not found, unable to create {env_name}:\"\n                f\" {task_name}. Consider downloading and installing \"\n                f\"dm_control from {self.git_url}\"\n            )\n\n        if _seed is not None:\n            random_state = np.random.RandomState(_seed)\n            kwargs = {\"random\": random_state}\n        env = suite.load(env_name, task_name, task_kwargs=kwargs)\n        return super()._build_env(\n            env, from_pixels=from_pixels, pixels_only=pixels_only, **kwargs\n        )\n\n    def rebuild_with_kwargs(self, **new_kwargs):\n        self._constructor_kwargs.update(new_kwargs)\n        self._env = self._build_env()\n        self._make_specs(self._env)\n\n    def _check_kwargs(self, kwargs: Dict):\n        if \"env_name\" in kwargs:\n            env_name = kwargs[\"env_name\"]\n            if \"task_name\" in kwargs:\n                task_name = kwargs[\"task_name\"]\n                if (\n                    env_name not in self.available_envs\n                    or task_name not in self.available_envs[env_name]\n                ):\n                    raise RuntimeError(\n                        f\"{env_name} with task {task_name} is unknown in {self.libname}\"\n                    )\n            else:\n                raise TypeError(\"dm_control requires task_name to be specified\")\n        else:\n            raise TypeError(\"dm_control requires env_name to be specified\")\n\n    # def _set_seed(self, _seed: int) -> int:\n    #     self._env = self._build_env(\n    #         _seed=_seed, **self._constructor_kwargs\n    #     )\n    #     self.reset()\n    #     return _seed\n\n    def __repr__(self) -> str:\n        return f\"{self.__class__.__name__}(env={self.env_name}, task={self.task_name}, batch_size={self.batch_size})\"\n", "path": "torchrl/envs/libs/dm_control.py"}], "after_files": [{"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\nfrom __future__ import annotations\n\nimport collections\nimport os\nfrom typing import Optional, Tuple, Union, Dict, Any\n\nimport numpy as np\nimport torch\n\nfrom torchrl.data import (\n    CompositeSpec,\n    NdBoundedTensorSpec,\n    NdUnboundedContinuousTensorSpec,\n    TensorSpec,\n    NdUnboundedDiscreteTensorSpec,\n)\nfrom ...data.utils import numpy_to_torch_dtype_dict, DEVICE_TYPING\nfrom ..gym_like import GymLikeEnv\n\nif torch.has_cuda and torch.cuda.device_count() > 1:\n    n = torch.cuda.device_count() - 1\n    os.environ[\"EGL_DEVICE_ID\"] = str(1 + (os.getpid() % n))\n    print(\"EGL_DEVICE_ID: \", os.environ[\"EGL_DEVICE_ID\"])\n\ntry:\n\n    import dm_control\n    import dm_env\n    from dm_control import suite\n    from dm_control.suite.wrappers import pixels\n\n    _has_dmc = True\n\nexcept ImportError:\n    _has_dmc = False\n\n__all__ = [\"DMControlEnv\", \"DMControlWrapper\"]\n\n\ndef _dmcontrol_to_torchrl_spec_transform(\n    spec,\n    dtype: Optional[torch.dtype] = None,\n    device: DEVICE_TYPING = None,\n) -> TensorSpec:\n    if isinstance(spec, collections.OrderedDict):\n        spec = {\n            \"next_\" + k: _dmcontrol_to_torchrl_spec_transform(item, device=device)\n            for k, item in spec.items()\n        }\n        return CompositeSpec(**spec)\n    elif isinstance(spec, dm_env.specs.BoundedArray):\n        if dtype is None:\n            dtype = numpy_to_torch_dtype_dict[spec.dtype]\n        return NdBoundedTensorSpec(\n            shape=spec.shape,\n            minimum=spec.minimum,\n            maximum=spec.maximum,\n            dtype=dtype,\n            device=device,\n        )\n    elif isinstance(spec, dm_env.specs.Array):\n        if dtype is None:\n            dtype = numpy_to_torch_dtype_dict[spec.dtype]\n        if dtype in (torch.float, torch.double, torch.half):\n            return NdUnboundedContinuousTensorSpec(\n                shape=spec.shape, dtype=dtype, device=device\n            )\n        else:\n            return NdUnboundedDiscreteTensorSpec(\n                shape=spec.shape, dtype=dtype, device=device\n            )\n\n    else:\n        raise NotImplementedError(type(spec))\n\n\ndef _get_envs(to_dict: bool = True) -> Dict[str, Any]:\n    if not _has_dmc:\n        return dict()\n    if not to_dict:\n        return tuple(suite.BENCHMARKING) + tuple(suite.EXTRA)\n    d = dict()\n    for tup in suite.BENCHMARKING:\n        env_name = tup[0]\n        d.setdefault(env_name, []).append(tup[1])\n    for tup in suite.EXTRA:\n        env_name = tup[0]\n        d.setdefault(env_name, []).append(tup[1])\n    return d\n\n\ndef _robust_to_tensor(array: Union[float, np.ndarray]) -> torch.Tensor:\n    if isinstance(array, np.ndarray):\n        return torch.tensor(array.copy())\n    else:\n        return torch.tensor(array)\n\n\nclass DMControlWrapper(GymLikeEnv):\n    \"\"\"\n    DeepMind Control lab environment wrapper.\n\n    Args:\n        env (dm_control.suite env): environment instance\n        from_pixels (bool): if True, the observation\n\n    Examples:\n        >>> env = dm_control.suite.load(\"cheetah\", \"run\")\n        >>> env = DMControlWrapper(env,\n        ...    from_pixels=True, frame_skip=4)\n        >>> td = env.rand_step()\n        >>> print(td)\n        >>> print(env.available_envs)\n    \"\"\"\n\n    git_url = \"https://github.com/deepmind/dm_control\"\n    libname = \"dm_control\"\n    available_envs = _get_envs()\n\n    def __init__(self, env=None, **kwargs):\n        if env is not None:\n            kwargs[\"env\"] = env\n        super().__init__(**kwargs)\n\n    def _build_env(\n        self,\n        env,\n        _seed: Optional[int] = None,\n        from_pixels: bool = False,\n        render_kwargs: Optional[dict] = None,\n        pixels_only: bool = False,\n        **kwargs,\n    ):\n        self.from_pixels = from_pixels\n        self.pixels_only = pixels_only\n\n        if from_pixels:\n            self._set_egl_device(self.device)\n            self.render_kwargs = {\"camera_id\": 0}\n            if render_kwargs is not None:\n                self.render_kwargs.update(render_kwargs)\n            env = pixels.Wrapper(\n                env,\n                pixels_only=self.pixels_only,\n                render_kwargs=self.render_kwargs,\n            )\n        return env\n\n    def _make_specs(self, env: \"gym.Env\") -> None:  # noqa: F821\n        # specs are defined when first called\n        return\n\n    def _check_kwargs(self, kwargs: Dict):\n        if \"env\" not in kwargs:\n            raise TypeError(\"Could not find environment key 'env' in kwargs.\")\n        env = kwargs[\"env\"]\n        if not isinstance(env, (dm_control.rl.control.Environment, pixels.Wrapper)):\n            raise TypeError(\n                \"env is not of type 'dm_control.rl.control.Environment' or `dm_control.suite.wrappers.pixels.Wrapper`.\"\n            )\n\n    def _set_egl_device(self, device: DEVICE_TYPING):\n        # Deprecated as lead to unreliable rendering\n        # egl device needs to be set before importing mujoco bindings: in\n        # distributed settings, it'll be easy to tell which cuda device to use.\n        # In mp settings, we'll need to use mp.Pool with a specific init function\n        # that defines the EGL device before importing libraries. For now, we'll\n        # just use a common EGL_DEVICE_ID environment variable for all processes.\n        return\n\n    def to(self, device: DEVICE_TYPING) -> DMControlEnv:\n        super().to(device)\n        self._set_egl_device(self.device)\n        return self\n\n    def _init_env(self, seed: Optional[int] = None) -> Optional[int]:\n        seed = self.set_seed(seed)\n        return seed\n\n    def _set_seed(self, _seed: Optional[int]) -> Optional[int]:\n        if _seed is None:\n            return None\n        random_state = np.random.RandomState(_seed)\n        if isinstance(self._env, pixels.Wrapper):\n            if not hasattr(self._env._env.task, \"_random\"):\n                raise RuntimeError(\"self._env._env.task._random does not exist\")\n            self._env._env.task._random = random_state\n        else:\n            if not hasattr(self._env.task, \"_random\"):\n                raise RuntimeError(\"self._env._env.task._random does not exist\")\n            self._env.task._random = random_state\n        self.reset()\n        return _seed\n\n    def _output_transform(\n        self, timestep_tuple: Tuple[\"TimeStep\"]  # noqa: F821\n    ) -> Tuple[np.ndarray, float, bool]:\n        if type(timestep_tuple) is not tuple:\n            timestep_tuple = (timestep_tuple,)\n        reward = timestep_tuple[0].reward\n\n        done = False  # dm_control envs are non-terminating\n        observation = timestep_tuple[0].observation\n        return observation, reward, done\n\n    @property\n    def input_spec(self) -> TensorSpec:\n        if self._input_spec is None:\n            self._input_spec = CompositeSpec(\n                action=_dmcontrol_to_torchrl_spec_transform(\n                    self._env.action_spec(), device=self.device\n                )\n            )\n        return self._input_spec\n\n    @input_spec.setter\n    def input_spec(self, value: TensorSpec) -> None:\n        self._input_spec = value\n\n    @property\n    def observation_spec(self) -> TensorSpec:\n        if self._observation_spec is None:\n            self._observation_spec = _dmcontrol_to_torchrl_spec_transform(\n                self._env.observation_spec(), device=self.device\n            )\n        return self._observation_spec\n\n    @observation_spec.setter\n    def observation_spec(self, value: TensorSpec) -> None:\n        self._observation_spec = value\n\n    @property\n    def reward_spec(self) -> TensorSpec:\n        if self._reward_spec is None:\n            self._reward_spec = _dmcontrol_to_torchrl_spec_transform(\n                self._env.reward_spec(), device=self.device\n            )\n        return self._reward_spec\n\n    @reward_spec.setter\n    def reward_spec(self, value: TensorSpec) -> None:\n        self._reward_spec = value\n\n    def __repr__(self) -> str:\n        return (\n            f\"{self.__class__.__name__}(env={self._env}, batch_size={self.batch_size})\"\n        )\n\n\nclass DMControlEnv(DMControlWrapper):\n    \"\"\"\n    DeepMind Control lab environment wrapper.\n\n    Args:\n        env_name (str): name of the environment\n        task_name (str): name of the task\n        seed (int, optional): seed to use for the environment\n        from_pixels (bool, optional): if True, the observation will be returned\n            as an image.\n            Default is False.\n\n    Examples:\n        >>> env = DMControlEnv(env_name=\"cheetah\", task_name=\"run\",\n        ...    from_pixels=True, frame_skip=4)\n        >>> td = env.rand_step()\n        >>> print(td)\n        >>> print(env.available_envs)\n    \"\"\"\n\n    def __init__(self, env_name, task_name, **kwargs):\n        if not _has_dmc:\n            raise ImportError(\n                \"dm_control python package was not found.\"\n                \"Please install this dependency.\"\n            )\n        kwargs[\"env_name\"] = env_name\n        kwargs[\"task_name\"] = task_name\n        super().__init__(**kwargs)\n\n    def _build_env(\n        self,\n        env_name: str,\n        task_name: str,\n        _seed: Optional[int] = None,\n        **kwargs,\n    ):\n        self.env_name = env_name\n        self.task_name = task_name\n\n        from_pixels = kwargs.get(\"from_pixels\")\n        if \"from_pixels\" in kwargs:\n            del kwargs[\"from_pixels\"]\n        pixels_only = kwargs.get(\"pixels_only\")\n        if \"pixels_only\" in kwargs:\n            del kwargs[\"pixels_only\"]\n\n        if not _has_dmc:\n            raise RuntimeError(\n                f\"dm_control not found, unable to create {env_name}:\"\n                f\" {task_name}. Consider downloading and installing \"\n                f\"dm_control from {self.git_url}\"\n            )\n\n        if _seed is not None:\n            random_state = np.random.RandomState(_seed)\n            kwargs = {\"random\": random_state}\n        env = suite.load(env_name, task_name, task_kwargs=kwargs)\n        return super()._build_env(\n            env, from_pixels=from_pixels, pixels_only=pixels_only, **kwargs\n        )\n\n    def rebuild_with_kwargs(self, **new_kwargs):\n        self._constructor_kwargs.update(new_kwargs)\n        self._env = self._build_env()\n        self._make_specs(self._env)\n\n    def _check_kwargs(self, kwargs: Dict):\n        if \"env_name\" in kwargs:\n            env_name = kwargs[\"env_name\"]\n            if \"task_name\" in kwargs:\n                task_name = kwargs[\"task_name\"]\n                if (\n                    env_name not in self.available_envs\n                    or task_name not in self.available_envs[env_name]\n                ):\n                    raise RuntimeError(\n                        f\"{env_name} with task {task_name} is unknown in {self.libname}\"\n                    )\n            else:\n                raise TypeError(\"dm_control requires task_name to be specified\")\n        else:\n            raise TypeError(\"dm_control requires env_name to be specified\")\n\n    # def _set_seed(self, _seed: int) -> int:\n    #     self._env = self._build_env(\n    #         _seed=_seed, **self._constructor_kwargs\n    #     )\n    #     self.reset()\n    #     return _seed\n\n    def __repr__(self) -> str:\n        return f\"{self.__class__.__name__}(env={self.env_name}, task={self.task_name}, batch_size={self.batch_size})\"\n", "path": "torchrl/envs/libs/dm_control.py"}]}