Merge remote-tracking branch 'origin/main' into fix-jumanji

pytorch · Apr 30, 2024 · 9210872 · 9210872
2 parents 46c0ce1 + 68101b0
commit 9210872
Show file tree

Hide file tree

Showing 6 changed files with 112 additions and 12 deletions.
diff --git a/.github/scripts/m1_script.sh b/.github/scripts/m1_script.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
 
-export BUILD_VERSION=0.4.0
+export TORCHRL_BUILD_VERSION=0.4.0
 
 ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
@@ -32,7 +32,7 @@ jobs:
         run: |
           export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH"
           python3 -mpip install wheel
-          BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel
+          TORCHRL_BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel
           # NB: wheels have the linux_x86_64 tag so we rename to manylinux1
           # find . -name 'dist/*whl' -exec bash -c ' mv $0 ${0/linux/manylinux1}' {} \;
       # pytorch/pytorch binaries are also manylinux_2_17 compliant but they
@@ -72,7 +72,7 @@ jobs:
         shell: bash
         run: |
             python3 -mpip install wheel
-            BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel
+            TORCHRL_BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel
       - name: Upload wheel for the test-wheel job
         uses: actions/upload-artifact@v2
         with:

diff --git a/examples/envs/gym-async-info-reader.py b/examples/envs/gym-async-info-reader.py
@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+A toy example of executing a Gym environment asynchronously and gathering the info properly.
+"""
+import argparse
+
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--use_wrapper", action="store_true")
+
+# Create the dummy environment
+class CustomEnv(gym.Env):
+    def __init__(self, render_mode=None):
+        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(3,))
+        self.action_space = spaces.Box(low=-np.inf, high=np.inf, shape=(1,))
+
+    def _get_info(self):
+        return {"field1": self.state**2}
+
+    def _get_obs(self):
+        return self.state.copy()
+
+    def reset(self, seed=None, options=None):
+        # We need the following line to seed self.np_random
+        super().reset(seed=seed)
+        self.state = np.zeros(self.observation_space.shape)
+        observation = self._get_obs()
+        info = self._get_info()
+        return observation, info
+
+    def step(self, action):
+        self.state += action.item()
+        truncated = False
+        terminated = False
+        reward = 1 if terminated else 0  # Binary sparse rewards
+        observation = self._get_obs()
+        info = self._get_info()
+        return observation, reward, terminated, truncated, info
+
+
+if __name__ == "__main__":
+    import torch
+    from torchrl.data.tensor_specs import UnboundedContinuousTensorSpec
+    from torchrl.envs import check_env_specs, GymEnv, GymWrapper
+
+    args = parser.parse_args()
+
+    num_envs = 10
+
+    if args.use_wrapper:
+        # Option 1: using GymWrapper
+        env = gym.vector.AsyncVectorEnv([lambda: CustomEnv() for _ in range(num_envs)])
+        env = GymWrapper(env, device="cpu")
+    else:
+        # Option 2: using GymEnv directly, no need to call AsyncVectorEnv
+        gym.register("Custom-v0", CustomEnv)
+        env = GymEnv("Custom-v0", num_envs=num_envs)
+
+    keys = ["field1"]
+    specs = [
+        UnboundedContinuousTensorSpec(shape=(num_envs, 3), dtype=torch.float64),
+    ]
+
+    # Create an info reader: this object will read the info and write its content to the tensordict
+    def reader(info, tensordict):
+        return tensordict.set("field1", np.stack(info["field1"]))
+    env.set_info_dict_reader(info_dict_reader=reader)
+
+    # Print the info readers (there should be 2: one to read the terminal states and another to read the 'field1')
+    print("readers", env.info_dict_reader)
+
+    # We need to unlock the specs to make them writable
+    env.observation_spec.unlock_()
+    env.observation_spec["field1"] = specs[0]
+    env.observation_spec.lock_()
+
+    # Check that we did a good job
+    check_env_specs(env)
+
+    td = env.reset()
+    print("reset data", td)
+    print("content of field1 (should be a 10x3 tensor)", td["field1"])
diff --git a/setup.py b/setup.py
@@ -32,8 +32,8 @@ def get_version():
     version_txt = os.path.join(cwd, "version.txt")
     with open(version_txt, "r") as f:
         version = f.readline().strip()
-    if os.getenv("BUILD_VERSION"):
-        version = os.getenv("BUILD_VERSION")
+    if os.getenv("TORCHRL_BUILD_VERSION"):
+        version = os.getenv("TORCHRL_BUILD_VERSION")
     elif sha != "Unknown":
         version += "+" + sha[:7]
     return version
@@ -68,11 +68,13 @@ def write_version_file(version):
         f.write("git_version = {}\n".format(repr(sha)))
 
 
-def _get_pytorch_version(is_nightly):
+def _get_pytorch_version(is_nightly, is_local):
     # if "PYTORCH_VERSION" in os.environ:
     #     return f"torch=={os.environ['PYTORCH_VERSION']}"
     if is_nightly:
         return "torch>=2.4.0.dev"
+    elif is_local:
+        return "torch"
     return "torch>=2.3.0"
 
 
@@ -178,10 +180,12 @@ def _main(argv):
     else:
         version = get_version()
         write_version_file(version)
+    TORCHRL_BUILD_VERSION = os.getenv("TORCHRL_BUILD_VERSION")
     logging.info("Building wheel {}-{}".format(package_name, version))
-    logging.info(f"BUILD_VERSION is {os.getenv('BUILD_VERSION')}")
+    logging.info(f"TORCHRL_BUILD_VERSION is {TORCHRL_BUILD_VERSION}")
 
-    pytorch_package_dep = _get_pytorch_version(is_nightly)
+    is_local = TORCHRL_BUILD_VERSION is None
+    pytorch_package_dep = _get_pytorch_version(is_nightly, is_local)
     logging.info("-- PyTorch dependency:", pytorch_package_dep)
     # branch = _run_cmd(["git", "rev-parse", "--abbrev-ref", "HEAD"])
     # tag = _run_cmd(["git", "describe", "--tags", "--exact-match", "@"])

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -1597,12 +1597,18 @@ def __init__(
             if high is not None:
                 raise TypeError(self.CONFLICTING_KWARGS.format("high", "maximum"))
             high = kwargs.pop("maximum")
-            warnings.warn(self.DEPRECATED_KWARGS, category=DeprecationWarning)
+            warnings.warn(
+                "Maximum is deprecated since v0.4.0, using high instead.",
+                category=DeprecationWarning,
+            )
         if "minimum" in kwargs:
             if low is not None:
                 raise TypeError(self.CONFLICTING_KWARGS.format("low", "minimum"))
             low = kwargs.pop("minimum")
-            warnings.warn(self.DEPRECATED_KWARGS, category=DeprecationWarning)
+            warnings.warn(
+                "Minimum is deprecated since v0.4.0, using low instead.",
+                category=DeprecationWarning,
+            )
         domain = kwargs.pop("domain", "continuous")
         if len(kwargs):
             raise TypeError(f"Got unrecognised kwargs {tuple(kwargs.keys())}.")

diff --git a/torchrl/envs/gym_like.py b/torchrl/envs/gym_like.py
@@ -162,8 +162,9 @@ class GymLikeEnv(_EnvWrapper):
 
     @classmethod
     def __new__(cls, *args, **kwargs):
-        cls._info_dict_reader = []
-        return super().__new__(cls, *args, _batch_locked=True, **kwargs)
+        self = super().__new__(cls, *args, _batch_locked=True, **kwargs)
+        self._info_dict_reader = []
+        return self
 
     def read_action(self, action):
         """Reads the action obtained from the input TensorDict and transforms it in the format expected by the contained environment.