[Tests] Remove _utils_internal.py in tests

test/smoke_test_deps.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import argparse
-import os
 import sys
 import tempfile
 
@@ -65,10 +64,7 @@ def test_gym():
         import ale_py  # noqa: F401
     except Exception:  # pragma: no cover
         pytest.skip("ALE not available (missing ale_py); skipping Atari gym test.")
-    if os.getenv("PYTORCH_TEST_FBCODE"):
-        from pytorch.rl.test._utils_internal import PONG_VERSIONED
-    else:
-        from _utils_internal import PONG_VERSIONED
+    from torchrl.testing import PONG_VERSIONED
 
     try:
         env = GymEnv(PONG_VERSIONED())

test/test_actors.py

@@ -6,7 +6,6 @@
 
 import argparse
 import importlib.util
-import os
 
 import pytest
 import torch
@@ -33,10 +32,7 @@
     ValueOperator,
 )
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import get_default_devices
-else:
-    from _utils_internal import get_default_devices
+from torchrl.testing import get_default_devices
 from torchrl.testing.mocking_classes import NestedCountingEnv
 
 _has_vllm = importlib.util.find_spec("vllm") is not None

test/test_collectors.py

@@ -8,7 +8,6 @@
 import contextlib
 import functools
 import gc
-import os
 import subprocess
 import sys
 import time
@@ -88,37 +87,17 @@
     RandomPolicy,
     SafeModule,
 )
-from torchrl.testing.modules import BiasModule, NonSerializableBiasModule
-from torchrl.testing.mp_helpers import decorate_thread_sub_func
-from torchrl.weight_update import (
-    MultiProcessWeightSyncScheme,
-    SharedMemWeightSyncScheme,
-)
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    IS_FB = True
-    from pytorch.rl.test._utils_internal import (
-        CARTPOLE_VERSIONED,
-        check_rollout_consistency_multikey_env,
-        generate_seeds,
-        get_available_devices,
-        get_default_devices,
-        LSTMNet,
-        PENDULUM_VERSIONED,
-        retry,
-    )
-else:
-    IS_FB = False
-    from _utils_internal import (
-        CARTPOLE_VERSIONED,
-        check_rollout_consistency_multikey_env,
-        generate_seeds,
-        get_available_devices,
-        get_default_devices,
-        LSTMNet,
-        PENDULUM_VERSIONED,
-        retry,
-    )
+from torchrl.testing import (
+    CARTPOLE_VERSIONED,
+    check_rollout_consistency_multikey_env,
+    generate_seeds,
+    get_available_devices,
+    get_default_devices,
+    LSTMNet,
+    PENDULUM_VERSIONED,
+    retry,
+)
 from torchrl.testing.mocking_classes import (
     ContinuousActionVecMockEnv,
     CountingBatchedEnv,
@@ -137,6 +116,12 @@
     MultiKeyCountingEnvPolicy,
     NestedCountingEnv,
 )
+from torchrl.testing.modules import BiasModule, NonSerializableBiasModule
+from torchrl.testing.mp_helpers import decorate_thread_sub_func
+from torchrl.weight_update import (
+    MultiProcessWeightSyncScheme,
+    SharedMemWeightSyncScheme,
+)
 
 # torch.set_default_dtype(torch.double)
 IS_WINDOWS = sys.platform == "win32"
@@ -864,7 +849,6 @@ def test_env_that_errors(self, ctype):
                 break
 
     @retry(AssertionError, tries=10, delay=0)
-    @pytest.mark.skipif(IS_FB, reason="Not compatible with fbcode")
     @pytest.mark.parametrize("to", [3, 10])
     @pytest.mark.parametrize(
         "collector_cls", ["MultiSyncCollector", "MultiAsyncCollector"]

test/test_distributions.py

@@ -6,7 +6,6 @@
 
 import argparse
 import importlib.util
-import os
 
 import pytest
 import torch
@@ -36,10 +35,7 @@
     LLMMaskedCategorical,
 )
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import get_default_devices
-else:
-    from _utils_internal import get_default_devices
+from torchrl.testing import get_default_devices
 
 _has_scipy = importlib.util.find_spec("scipy", None) is not None
 

test/test_envs.py

@@ -174,28 +174,16 @@ def check_no_lingering_multiprocessing_resources(request):
     _atari_found = False
     atari_confs = defaultdict(str)
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import (
-        _make_envs,
-        CARTPOLE_VERSIONED,
-        check_rollout_consistency_multikey_env,
-        get_default_devices,
-        HALFCHEETAH_VERSIONED,
-        PENDULUM_VERSIONED,
-        PONG_VERSIONED,
-        rand_reset,
-    )
-else:
-    from _utils_internal import (
-        _make_envs,
-        CARTPOLE_VERSIONED,
-        check_rollout_consistency_multikey_env,
-        get_default_devices,
-        HALFCHEETAH_VERSIONED,
-        PENDULUM_VERSIONED,
-        PONG_VERSIONED,
-        rand_reset,
-    )
+from torchrl.testing import (
+    CARTPOLE_VERSIONED,
+    check_rollout_consistency_multikey_env,
+    get_default_devices,
+    HALFCHEETAH_VERSIONED,
+    make_envs as _make_envs,
+    PENDULUM_VERSIONED,
+    PONG_VERSIONED,
+    rand_reset,
+)
 from torchrl.testing.mocking_classes import (
     ActionObsMergeLinear,
     AutoResetHeteroCountingEnv,

test/test_exploration.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import argparse
-import os
 
 import pytest
 import torch
@@ -41,10 +40,7 @@
     OrnsteinUhlenbeckProcessModule,
 )
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import get_default_devices
-else:
-    from _utils_internal import get_default_devices
+from torchrl.testing import get_default_devices
 from torchrl.testing.mocking_classes import (
     ContinuousActionVecMockEnv,
     CountingEnvCountModule,

test/test_helpers.py

@@ -6,7 +6,6 @@
 
 import argparse
 import dataclasses
-import os
 import pathlib
 import sys
 from time import sleep
@@ -24,6 +23,15 @@
     FlattenObservation,
     TransformedEnv,
 )
+
+from torchrl.testing import generate_seeds, get_default_devices
+from torchrl.testing.mocking_classes import (
+    ContinuousActionConvMockEnvNumpy,
+    ContinuousActionVecMockEnv,
+    DiscreteActionConvMockEnvNumpy,
+    DiscreteActionVecMockEnv,
+    MockSerialEnv,
+)
 from torchrl.trainers.helpers import transformed_env_constructor
 from torchrl.trainers.helpers.envs import (
     EnvConfig,
@@ -36,18 +44,6 @@
     make_dqn_actor,
 )
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import generate_seeds, get_default_devices
-else:
-    from _utils_internal import generate_seeds, get_default_devices
-from torchrl.testing.mocking_classes import (
-    ContinuousActionConvMockEnvNumpy,
-    ContinuousActionVecMockEnv,
-    DiscreteActionConvMockEnvNumpy,
-    DiscreteActionVecMockEnv,
-    MockSerialEnv,
-)
-
 try:
     from hydra import compose, initialize
     from hydra.core.config_store import ConfigStore

test/test_libs.py

@@ -136,34 +136,19 @@
 _has_ray = importlib.util.find_spec("ray") is not None
 _has_ale = importlib.util.find_spec("ale_py") is not None
 _has_mujoco = importlib.util.find_spec("mujoco") is not None
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import (
-        _make_multithreaded_env,
-        CARTPOLE_VERSIONED,
-        CLIFFWALKING_VERSIONED,
-        get_available_devices,
-        get_default_devices,
-        HALFCHEETAH_VERSIONED,
-        PENDULUM_VERSIONED,
-        PONG_VERSIONED,
-        rand_reset,
-        retry,
-        rollout_consistency_assertion,
-    )
-else:
-    from _utils_internal import (
-        _make_multithreaded_env,
-        CARTPOLE_VERSIONED,
-        CLIFFWALKING_VERSIONED,
-        get_available_devices,
-        get_default_devices,
-        HALFCHEETAH_VERSIONED,
-        PENDULUM_VERSIONED,
-        PONG_VERSIONED,
-        rand_reset,
-        retry,
-        rollout_consistency_assertion,
-    )
+from torchrl.testing import (
+    CARTPOLE_VERSIONED,
+    CLIFFWALKING_VERSIONED,
+    get_available_devices,
+    get_default_devices,
+    HALFCHEETAH_VERSIONED,
+    make_multithreaded_env as _make_multithreaded_env,
+    PENDULUM_VERSIONED,
+    PONG_VERSIONED,
+    rand_reset,
+    retry,
+    rollout_consistency_assertion,
+)
 
 TORCH_VERSION = version.parse(version.parse(torch.__version__).base_version)
 

test/test_modules.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import argparse
-import os
 import re
 
 from numbers import Number
@@ -58,12 +57,9 @@
 from torchrl.modules.planners.mppi import MPPIPlanner
 from torchrl.objectives.value import TDLambdaEstimator
 
-from torchrl.testing.mocking_classes import MockBatchedUnLockedEnv
+from torchrl.testing import get_default_devices, retry
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import get_default_devices, retry
-else:
-    from _utils_internal import get_default_devices, retry
+from torchrl.testing.mocking_classes import MockBatchedUnLockedEnv
 
 
 @pytest.fixture

test/test_objectives.py

@@ -10,7 +10,6 @@
 import importlib.util
 import itertools
 import operator
-import os
 import sys
 import warnings
 from copy import deepcopy
@@ -155,22 +154,13 @@
     _split_and_pad_sequence,
 )
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import (  # noqa
-        _call_value_nets,
-        dtype_fixture,
-        get_available_devices,
-        get_default_devices,
-        PENDULUM_VERSIONED,
-    )
-else:
-    from _utils_internal import (  # noqa
-        _call_value_nets,
-        dtype_fixture,
-        get_available_devices,
-        get_default_devices,
-        PENDULUM_VERSIONED,
-    )
+from torchrl.testing import (  # noqa
+    call_value_nets as _call_value_nets,
+    dtype_fixture,
+    get_available_devices,
+    get_default_devices,
+    PENDULUM_VERSIONED,
+)
 from torchrl.testing.mocking_classes import ContinuousActionConvMockEnv
 
 _has_functorch = True

test/test_postprocs.py

@@ -6,7 +6,6 @@
 
 import argparse
 import functools
-import os
 
 import pytest
 import torch
@@ -16,10 +15,7 @@
 from torchrl.collectors.utils import split_trajectories
 from torchrl.data.postprocs.postprocs import DensifyReward, MultiStep
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import get_default_devices
-else:
-    from _utils_internal import get_default_devices
+from torchrl.testing import get_default_devices
 
 
 class TestMultiStep: