[Refactor] Import enum and interaction_type utils (pytorch#1055)

ai4co · Apr 13, 2023 · d9bdac9 · d9bdac9
1 parent 40004e6
commit d9bdac9
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -440,9 +440,9 @@ And it is `functorch` and `torch.compile` compatible!
 
   ```python
   policy_explore = EGreedyWrapper(policy)
-  with set_exploration_mode("random"):
+  with set_exploration_type(ExplorationType.RANDOM):
       tensordict = policy_explore(tensordict)  # will use eps-greedy
-  with set_exploration_mode("mode"):
+  with set_exploration_type(ExplorationType.MODE):
       tensordict = policy_explore(tensordict)  # will not use eps-greedy
   ```
   </details>

diff --git a/docs/source/reference/envs.rst b/docs/source/reference/envs.rst
@@ -372,8 +372,10 @@ Helpers
 
     step_mdp
     get_available_libraries
-    set_exploration_mode
-    exploration_mode
+    set_exploration_mode #deprecated
+    set_exploration_type
+    exploration_mode #deprecated
+    exploration_type
     check_env_specs
     make_composite_from_td
 

diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
@@ -7,8 +7,14 @@
 import pkg_resources
 import torch
 from tensordict.nn.probabilistic import (  # noqa
+    # Note: the `set_/interaction_mode` and their associated arg `default_interaction_mode` are being deprecated!
+    #       Please use the `set_/interaction_type` ones above with the InteractionType enum instead.
+    #       See more details: https://github.com/pytorch/rl/issues/1016
     interaction_mode as exploration_mode,
+    interaction_type as exploration_type,
+    InteractionType as ExplorationType,
     set_interaction_mode as set_exploration_mode,
+    set_interaction_type as set_exploration_type,
 )
 from tensordict.tensordict import TensorDictBase