torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/testing/assertions.py

@@ -0,0 +1,179 @@
+# 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.
+
+"""Assertions and validation utilities for TorchRL tests."""
+
+from __future__ import annotations
+
+import torch
+from tensordict import TensorDict
+
+__all__ = [
+    "check_rollout_consistency_multikey_env",
+    "rand_reset",
+    "rollout_consistency_assertion",
+]
+
+
+def rollout_consistency_assertion(
+    rollout, *, done_key="done", observation_key="observation", done_strict=False
+):
+    """Test that observations in 'next' match observations in the next root tensordict.
+
+    Verifies consistency: when done is False the next observation should match,
+    and when done is True they should differ (indicating a reset occurred).
+
+    Args:
+        rollout: The rollout tensordict to validate.
+        done_key: The key for the done signal.
+        observation_key: The key for observations.
+        done_strict: If True, raise an error if no done is detected.
+    """
+    done = rollout[..., :-1]["next", done_key].squeeze(-1)
+    # data resulting from step, when it's not done
+    r_not_done = rollout[..., :-1]["next"][~done]
+    # data resulting from step, when it's not done, after step_mdp
+    r_not_done_tp1 = rollout[:, 1:][~done]
+    torch.testing.assert_close(
+        r_not_done[observation_key],
+        r_not_done_tp1[observation_key],
+        msg=f"Key {observation_key} did not match",
+    )
+
+    if done_strict and not done.any():
+        raise RuntimeError("No done detected, test could not complete.")
+    if done.any():
+        # data resulting from step, when it's done
+        r_done = rollout[..., :-1]["next"][done]
+        # data resulting from step, when it's done, after step_mdp and reset
+        r_done_tp1 = rollout[..., 1:][done]
+        # check that at least one obs after reset does not match the version before reset
+        assert not torch.isclose(
+            r_done[observation_key], r_done_tp1[observation_key]
+        ).all()
+
+
+def rand_reset(env):
+    """Generate a tensordict with reset keys that mimic the done spec.
+
+    Values are drawn at random until at least one reset is present.
+
+    Args:
+        env: The environment to generate reset keys for.
+
+    Returns:
+        A TensorDict containing the reset signals.
+    """
+    full_done_spec = env.full_done_spec
+    result = {}
+    for reset_key, list_of_done in zip(env.reset_keys, env.done_keys_groups):
+        val = full_done_spec[list_of_done[0]].rand()
+        while not val.any():
+            val = full_done_spec[list_of_done[0]].rand()
+        result[reset_key] = val
+    # create a data structure that keeps the batch size of the nested specs
+    result = (
+        full_done_spec.zero().update(result).exclude(*full_done_spec.keys(True, True))
+    )
+    return result
+
+
+def check_rollout_consistency_multikey_env(td: TensorDict, max_steps: int):
+    """Check rollout consistency for environments with multiple observation/action keys.
+
+    Validates that:
+    - Done and reset behavior is correct for root, nested_1, and nested_2
+    - Observations update correctly based on actions
+    - Rewards are computed correctly
+
+    Args:
+        td: The rollout tensordict to validate.
+        max_steps: The maximum steps before done in the environment.
+    """
+    index_batch_size = (0,) * (len(td.batch_size) - 1)
+
+    # Check done and reset for root
+    observation_is_max = td["next", "observation"][..., 0, 0, 0] == max_steps + 1
+    next_is_done = td["next", "done"][index_batch_size][:-1].squeeze(-1)
+    assert (td["next", "done"][observation_is_max]).all()
+    assert (~td["next", "done"][~observation_is_max]).all()
+    # Obs after done is 0
+    assert (td["observation"][index_batch_size][1:][next_is_done] == 0).all()
+    # Obs after not done is previous obs
+    assert (
+        td["observation"][index_batch_size][1:][~next_is_done]
+        == td["next", "observation"][index_batch_size][:-1][~next_is_done]
+    ).all()
+    # Check observation and reward update with count action for root
+    action_is_count = td["action"].long().argmax(-1).to(torch.bool)
+    assert (
+        td["next", "observation"][action_is_count]
+        == td["observation"][action_is_count] + 1
+    ).all()
+    assert (td["next", "reward"][action_is_count] == 1).all()
+    # Check observation and reward do not update with no-count action for root
+    assert (
+        td["next", "observation"][~action_is_count]
+        == td["observation"][~action_is_count]
+    ).all()
+    assert (td["next", "reward"][~action_is_count] == 0).all()
+
+    # Check done and reset for nested_1
+    observation_is_max = td["next", "nested_1", "observation"][..., 0] == max_steps + 1
+    # done at the root always prevail
+    next_is_done = td["next", "done"][index_batch_size][:-1].squeeze(-1)
+    assert (td["next", "nested_1", "done"][observation_is_max]).all()
+    assert (~td["next", "nested_1", "done"][~observation_is_max]).all()
+    # Obs after done is 0
+    assert (
+        td["nested_1", "observation"][index_batch_size][1:][next_is_done] == 0
+    ).all()
+    # Obs after not done is previous obs
+    assert (
+        td["nested_1", "observation"][index_batch_size][1:][~next_is_done]
+        == td["next", "nested_1", "observation"][index_batch_size][:-1][~next_is_done]
+    ).all()
+    # Check observation and reward update with count action for nested_1
+    action_is_count = td["nested_1"]["action"].to(torch.bool)
+    assert (
+        td["next", "nested_1", "observation"][action_is_count]
+        == td["nested_1", "observation"][action_is_count] + 1
+    ).all()
+    assert (td["next", "nested_1", "gift"][action_is_count] == 1).all()
+    # Check observation and reward do not update with no-count action for nested_1
+    assert (
+        td["next", "nested_1", "observation"][~action_is_count]
+        == td["nested_1", "observation"][~action_is_count]
+    ).all()
+    assert (td["next", "nested_1", "gift"][~action_is_count] == 0).all()
+
+    # Check done and reset for nested_2
+    observation_is_max = td["next", "nested_2", "observation"][..., 0] == max_steps + 1
+    # done at the root always prevail
+    next_is_done = td["next", "done"][index_batch_size][:-1].squeeze(-1)
+    assert (td["next", "nested_2", "done"][observation_is_max]).all()
+    assert (~td["next", "nested_2", "done"][~observation_is_max]).all()
+    # Obs after done is 0
+    assert (
+        td["nested_2", "observation"][index_batch_size][1:][next_is_done] == 0
+    ).all()
+    # Obs after not done is previous obs
+    assert (
+        td["nested_2", "observation"][index_batch_size][1:][~next_is_done]
+        == td["next", "nested_2", "observation"][index_batch_size][:-1][~next_is_done]
+    ).all()
+    # Check observation and reward update with count action for nested_2
+    action_is_count = td["nested_2"]["azione"].squeeze(-1).to(torch.bool)
+    assert (
+        td["next", "nested_2", "observation"][action_is_count]
+        == td["nested_2", "observation"][action_is_count] + 1
+    ).all()
+    assert (td["next", "nested_2", "reward"][action_is_count] == 1).all()
+    # Check observation and reward do not update with no-count action for nested_2
+    assert (
+        td["next", "nested_2", "observation"][~action_is_count]
+        == td["nested_2", "observation"][~action_is_count]
+    ).all()
+    assert (td["next", "nested_2", "reward"][~action_is_count] == 0).all()

torchrl/testing/dist_utils.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import os
+import time
+from collections.abc import Callable
+from typing import Any
+
+import psutil
+
+__all__ = [
+    "assert_no_new_python_processes",
+    "is_python_process",
+    "snapshot_python_processes",
+]
+
+
+def is_python_process(comm: str | None, args: str | None) -> bool:
+    """Check if a process is a python process."""
+    if comm is None:
+        comm = ""
+    comm = comm.lower()
+    if comm.startswith(("python", "pypy")):
+        return True
+    if not args:
+        return False
+    return "python" in args.lower()
+
+
+def snapshot_python_processes(
+    root: psutil.Process | None = None,
+) -> dict[tuple[int, float], dict[str, Any]]:
+    """Snapshot python processes belonging to the given process tree.
+
+    Returns a dict keyed by (pid, start_time) -> info.
+    """
+    if root is None:
+        root = psutil.Process(os.getpid())
+
+    uid = os.getuid()
+
+    # Snapshot descendant PIDs first, then query process info via process_iter.
+    # This avoids race conditions where a child exits between `children()` and
+    # attribute access on a stale Process handle (common with Ray helpers).
+    descendant_pids = {root.pid}
+    descendant_pids.update(p.pid for p in root.children(recursive=True))
+
+    out: dict[tuple[int, float], dict[str, Any]] = {}
+    for proc in psutil.process_iter(
+        attrs=["pid", "name", "cmdline", "create_time", "uids"], ad_value=None
+    ):
+        info = proc.info
+        pid = info.get("pid")
+        if pid is None or pid not in descendant_pids:
+            continue
+        uids = info.get("uids")
+        if uids is None or uids.real != uid:
+            continue
+
+        name = info.get("name") or ""
+        cmdline = info.get("cmdline") or []
+        args = " ".join(cmdline) if isinstance(cmdline, (list, tuple)) else str(cmdline)
+        if not is_python_process(name, args):
+            continue
+
+        start_time = float(info.get("create_time") or 0.0)
+        key = (int(pid), start_time)
+        out[key] = {
+            "pid": int(pid),
+            "start_time": start_time,
+            "comm": name,
+            "args": args,
+        }
+    return out
+
+
+def assert_no_new_python_processes(
+    *,
+    baseline: dict[tuple[int, float], dict[str, Any]],
+    baseline_time: float,
+    timeout: float = 20.0,
+    ignore_info_fn: Callable[[dict[str, Any]], bool] | None = None,
+) -> None:
+    """Assert that no python process started after baseline_time remains alive.
+
+    The check is limited to the current process tree (pytest process + descendants).
+    """
+    if ignore_info_fn is None:
+
+        def ignore_info_fn(_info: dict[str, Any]) -> bool:
+            return False
+
+    deadline = time.time() + timeout
+    last_new: dict[tuple[int, float], dict[str, Any]] | None = None
+    while time.time() < deadline:
+        current = snapshot_python_processes()
+        new: dict[tuple[int, float], dict[str, Any]] = {}
+        for (pid, start_time), info in current.items():
+            if pid == os.getpid():
+                continue
+            if ignore_info_fn(info):
+                continue
+            # Guard against pid reuse: only consider processes started after the baseline.
+            if start_time and start_time < baseline_time - 1.0:
+                continue
+            if (pid, start_time) in baseline:
+                continue
+            new[(pid, start_time)] = info
+        if not new:
+            return
+        last_new = new
+        time.sleep(0.25)
+
+    if last_new is None:
+        return
+    details = "\n".join(
+        f"- pid={v['pid']} comm={v.get('comm')} args={v.get('args')}"
+        for v in last_new.values()
+    )
+    raise AssertionError(
+        "Leaked python processes detected after collector.shutdown().\n"
+        f"Processes still alive:\n{details}"
+    )

torchrl/testing/env_creators.py

@@ -0,0 +1,227 @@
+# 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.
+
+"""Environment creation utilities for TorchRL tests."""
+
+from __future__ import annotations
+
+import torch
+
+from torchrl.envs import MultiThreadedEnv, ObservationNorm
+from torchrl.envs.batched_envs import ParallelEnv, SerialEnv
+from torchrl.envs.libs.envpool import _has_envpool
+from torchrl.envs.libs.gym import GymEnv
+from torchrl.envs.transforms import (
+    Compose,
+    RewardClipping,
+    ToTensorImage,
+    TransformedEnv,
+)
+from torchrl.testing.gym_helpers import HALFCHEETAH_VERSIONED, PONG_VERSIONED
+from torchrl.testing.utils import mp_ctx
+
+__all__ = [
+    "get_transform_out",
+    "make_envs",
+    "make_multithreaded_env",
+]
+
+
+def make_envs(
+    env_name,
+    frame_skip,
+    transformed_in,
+    transformed_out,
+    N,
+    device="cpu",
+    kwargs=None,
+    local_mp_ctx=mp_ctx,
+):
+    """Create parallel, serial, multithreaded, and single environment instances.
+
+    This helper creates environments suitable for testing batched environment behavior.
+
+    Args:
+        env_name: The gym environment name.
+        frame_skip: Number of frames to skip.
+        transformed_in: Whether to apply transforms inside the base env.
+        transformed_out: Whether to apply transforms outside the batched env.
+        N: Number of environments in the batch.
+        device: Device for the environments.
+        kwargs: Additional keyword arguments for environment creation.
+        local_mp_ctx: Multiprocessing context ('fork' or 'spawn').
+
+    Returns:
+        Tuple of (env_parallel, env_serial, env_multithread, env0).
+    """
+    torch.manual_seed(0)
+    if not transformed_in:
+
+        def create_env_fn():
+            return GymEnv(env_name, frame_skip=frame_skip, device=device)
+
+    else:
+        if env_name == PONG_VERSIONED():
+
+            def create_env_fn():
+                base_env = GymEnv(env_name, frame_skip=frame_skip, device=device)
+                in_keys = list(base_env.observation_spec.keys(True, True))[:1]
+                return TransformedEnv(
+                    base_env,
+                    Compose(*[ToTensorImage(in_keys=in_keys), RewardClipping(0, 0.1)]),
+                )
+
+        else:
+
+            def create_env_fn():
+
+                base_env = GymEnv(env_name, frame_skip=frame_skip, device=device)
+                in_keys = list(base_env.observation_spec.keys(True, True))[:1]
+
+                return TransformedEnv(
+                    base_env,
+                    Compose(
+                        ObservationNorm(in_keys=in_keys, loc=0.5, scale=1.1),
+                        RewardClipping(0, 0.1),
+                    ),
+                )
+
+    env0 = create_env_fn()
+    env_parallel = ParallelEnv(
+        N, create_env_fn, create_env_kwargs=kwargs, mp_start_method=local_mp_ctx
+    )
+    env_serial = SerialEnv(N, create_env_fn, create_env_kwargs=kwargs)
+
+    for key in env0.observation_spec.keys(True, True):
+        obs_key = key
+        break
+    else:
+        obs_key = None
+
+    if transformed_out:
+        t_out = get_transform_out(env_name, transformed_in, obs_key=obs_key)
+
+        env0 = TransformedEnv(
+            env0,
+            t_out(),
+        )
+        env_parallel = TransformedEnv(
+            env_parallel,
+            t_out(),
+        )
+        env_serial = TransformedEnv(
+            env_serial,
+            t_out(),
+        )
+    else:
+        t_out = None
+
+    if _has_envpool:
+        env_multithread = make_multithreaded_env(
+            env_name,
+            frame_skip,
+            t_out,
+            N,
+            device="cpu",
+            kwargs=None,
+        )
+    else:
+        env_multithread = None
+
+    return env_parallel, env_serial, env_multithread, env0
+
+
+def make_multithreaded_env(
+    env_name,
+    frame_skip,
+    transformed_out,
+    N,
+    device="cpu",
+    kwargs=None,
+):
+    """Create a multithreaded environment using envpool.
+
+    Args:
+        env_name: The gym environment name.
+        frame_skip: Number of frames to skip.
+        transformed_out: Transform factory to apply, or None.
+        N: Number of environments in the batch.
+        device: Device for the environment.
+        kwargs: Additional keyword arguments (unused, for API compatibility).
+
+    Returns:
+        A MultiThreadedEnv instance, optionally wrapped with transforms.
+    """
+    torch.manual_seed(0)
+    multithreaded_kwargs = (
+        {"frame_skip": frame_skip} if env_name == PONG_VERSIONED() else {}
+    )
+    env_multithread = MultiThreadedEnv(
+        N,
+        env_name,
+        create_env_kwargs=multithreaded_kwargs,
+        device=device,
+    )
+
+    if transformed_out:
+        for key in env_multithread.observation_spec.keys(True, True):
+            obs_key = key
+            break
+        else:
+            obs_key = None
+        env_multithread = TransformedEnv(
+            env_multithread,
+            get_transform_out(env_name, transformed_in=False, obs_key=obs_key)(),
+        )
+    return env_multithread
+
+
+def get_transform_out(env_name, transformed_in, obs_key=None):
+    """Create a transform factory for output transforms based on environment type.
+
+    Args:
+        env_name: The gym environment name.
+        transformed_in: Whether transforms were already applied inside.
+        obs_key: The observation key to transform.
+
+    Returns:
+        A callable that returns a Compose transform.
+    """
+    if env_name == PONG_VERSIONED():
+        if obs_key is None:
+            obs_key = "pixels"
+
+        def t_out():
+            return (
+                Compose(*[ToTensorImage(in_keys=[obs_key]), RewardClipping(0, 0.1)])
+                if not transformed_in
+                else Compose(*[ObservationNorm(in_keys=[obs_key], loc=0, scale=1)])
+            )
+
+    elif env_name == HALFCHEETAH_VERSIONED:
+        if obs_key is None:
+            obs_key = ("observation", "velocity")
+
+        def t_out():
+            return Compose(
+                ObservationNorm(in_keys=[obs_key], loc=0.5, scale=1.1),
+                RewardClipping(0, 0.1),
+            )
+
+    else:
+        if obs_key is None:
+            obs_key = "observation"
+
+        def t_out():
+            return (
+                Compose(
+                    ObservationNorm(in_keys=[obs_key], loc=0.5, scale=1.1),
+                    RewardClipping(0, 0.1),
+                )
+                if not transformed_in
+                else Compose(ObservationNorm(in_keys=[obs_key], loc=1.0, scale=1.0))
+            )
+
+    return t_out

torchrl/testing/env_helper.py

@@ -0,0 +1,35 @@
+# 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.
+
+from __future__ import annotations
+
+
+def make_isaac_env(env_name: str = "Isaac-Ant-v0"):
+    """Helper function to create an IsaacLab env."""
+    import torch
+
+    torch.manual_seed(0)
+    import argparse
+
+    # This code block ensures that the Isaac app is started in headless mode
+    from isaaclab.app import AppLauncher
+    from torchrl import logger as torchrl_logger
+
+    parser = argparse.ArgumentParser(description="Train an RL agent with TorchRL.")
+    AppLauncher.add_app_launcher_args(parser)
+    args_cli, hydra_args = parser.parse_known_args(["--headless"])
+    AppLauncher(args_cli)
+
+    # Imports and env
+    import gymnasium as gym
+    import isaaclab_tasks  # noqa: F401
+    from isaaclab_tasks.manager_based.classic.ant.ant_env_cfg import AntEnvCfg
+    from torchrl.envs.libs.isaac_lab import IsaacLabWrapper
+
+    torchrl_logger.info("Making IsaacLab env...")
+    env = gym.make(env_name, cfg=AntEnvCfg())
+    torchrl_logger.info("Wrapping IsaacLab env...")
+    env = IsaacLabWrapper(env)
+    return env