[BE] Add type annotation for tensor_keys to facilitate auto-complete

ghstack-source-id: b4a8fe3
Pull Request resolved: #2696

Author: Vincent Moens

torchrl/objectives/a2c.py

@@ -242,6 +242,7 @@ class _AcceptedKeys:
         terminated: NestedKey = "terminated"
         sample_log_prob: NestedKey = "sample_log_prob"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator: ValueEstimators = ValueEstimators.GAE
 

torchrl/objectives/common.py

@@ -128,6 +128,7 @@ class _AcceptedKeys:
 
         pass
 
+    tensor_keys: _AcceptedKeys
     _vmap_randomness = None
     default_value_estimator: ValueEstimators = None
 

torchrl/objectives/cql.py

@@ -260,6 +260,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
 
@@ -1024,6 +1025,7 @@ class _AcceptedKeys:
         terminated: NestedKey = "terminated"
         pred_val: NestedKey = "pred_val"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = [

torchrl/objectives/crossq.py

@@ -242,6 +242,7 @@ class _AcceptedKeys:
         terminated: NestedKey = "terminated"
         log_prob: NestedKey = "_log_prob"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
 

torchrl/objectives/ddpg.py

@@ -173,6 +173,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator: ValueEstimators = ValueEstimators.TD0
     out_keys = [

torchrl/objectives/decision_transformer.py

@@ -70,6 +70,7 @@ class _AcceptedKeys:
         # the "action" output from the model
         action_pred: NestedKey = "action"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
 
     actor_network: TensorDictModule
@@ -280,6 +281,7 @@ class _AcceptedKeys:
         # the "action" output from the model
         action_pred: NestedKey = "action"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
 
     actor_network: TensorDictModule

torchrl/objectives/deprecated.py

@@ -127,6 +127,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     delay_actor: bool = False
     default_value_estimator = ValueEstimators.TD0

torchrl/objectives/dqn.py

@@ -164,6 +164,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = ["loss"]
@@ -435,6 +436,7 @@ class _AcceptedKeys:
         terminated: NestedKey = "terminated"
         steps_to_next_obs: NestedKey = "steps_to_next_obs"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
 

torchrl/objectives/dreamer.py

@@ -89,8 +89,13 @@ class _AcceptedKeys:
         pixels: NestedKey = "pixels"
         reco_pixels: NestedKey = "reco_pixels"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
 
+    decoder: TensorDictModule
+    reward_model: TensorDictModule
+    world_mdel: TensorDictModule
+
     def __init__(
         self,
         world_model: TensorDictModule,
@@ -238,9 +243,13 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TDLambda
 
+    value_model: TensorDictModule
+    actor_model: TensorDictModule
+
     def __init__(
         self,
         actor_model: TensorDictModule,
@@ -392,8 +401,11 @@ class _AcceptedKeys:
 
         value: NestedKey = "state_value"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
 
+    value_model: TensorDictModule
+
     def __init__(
         self,
         value_model: TensorDictModule,

torchrl/objectives/gail.py

@@ -59,6 +59,7 @@ class _AcceptedKeys:
         collector_observation: NestedKey = "collector_observation"
         discriminator_pred: NestedKey = "d_logits"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
 
     discriminator_network: TensorDictModule

torchrl/objectives/iql.py

@@ -233,6 +233,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = [
@@ -709,6 +710,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = [

torchrl/objectives/multiagent/qmixer.py

@@ -179,6 +179,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = ["loss"]

torchrl/objectives/ppo.py

@@ -295,6 +295,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.GAE
 

torchrl/objectives/redq.py

@@ -231,6 +231,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     delay_actor: bool = False
     default_value_estimator = ValueEstimators.TD0

torchrl/objectives/reinforce.py

@@ -211,6 +211,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.GAE
     out_keys = ["loss_actor", "loss_value"]

torchrl/objectives/sac.py

@@ -290,6 +290,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
 
@@ -1029,6 +1030,7 @@ class _AcceptedKeys:
         terminated: NestedKey = "terminated"
         log_prob: NestedKey = "log_prob"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     delay_actor: bool = False

torchrl/objectives/td3.py

@@ -204,6 +204,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = [

torchrl/objectives/td3_bc.py

@@ -217,6 +217,7 @@ class _AcceptedKeys:
         done: NestedKey = "done"
         terminated: NestedKey = "terminated"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     default_value_estimator = ValueEstimators.TD0
     out_keys = [

torchrl/objectives/value/advantages.py

@@ -143,6 +143,7 @@ class _AcceptedKeys:
         steps_to_next_obs: NestedKey = "steps_to_next_obs"
         sample_log_prob: NestedKey = "sample_log_prob"
 
+    tensor_keys: _AcceptedKeys
     default_keys = _AcceptedKeys()
     value_network: Union[TensorDictModule, Callable]
     _vmap_randomness = None