Merge pull request #89 from UT-Austin-RPL/bugs

add option to pass keys from observation directly to actor net

Author: jakegrigsby

amago/agent.py

@@ -3,7 +3,7 @@
 """
 
 import itertools
-from typing import Type, Optional, Tuple, Any, List
+from typing import Type, Optional, Tuple, Any, List, Iterable
 
 import torch
 from torch import nn
@@ -222,6 +222,7 @@ def __init__(
         use_multigamma: bool = True,
         actor_type: Type[actor_critic.BaseActorHead] = actor_critic.Actor,
         critic_type: Type[actor_critic.BaseCriticHead] = actor_critic.NCritics,
+        pass_obs_keys_to_actor: Optional[Iterable[str]] = None,
     ):
         super().__init__()
         self.obs_space = obs_space
@@ -293,6 +294,7 @@ def __init__(
         self.target_actor = actor_type(**ac_kwargs)
         # full weight copy to targets
         self.hard_sync_targets()
+        self.pass_obs_keys_to_actor = pass_obs_keys_to_actor or []
 
     @property
     def trainable_params(self):
@@ -364,7 +366,10 @@ def get_actions(
             tstep_emb, time_idxs=time_idxs, hidden_state=hidden_state
         )
         # generate action distribution [batch, length, len(self.gammas), d_action]
-        action_dists = self.actor(traj_emb_t)
+        action_dists = self.actor(
+            traj_emb_t,
+            straight_from_obs={k: obs[k] for k in self.pass_obs_keys_to_actor},
+        )
         if sample:
             actions = action_dists.sample()
         else:
@@ -473,7 +478,7 @@ def forward(
         ## a ~ \pi(s) ##
         ################
         critic_loss = None
-        a_dist = self.actor(s_rep, log_dict=active_log_dict)
+        a_dist = self.actor(s_rep, log_dict=active_log_dict, straight_from_obs={k : batch.obs[k] for k in self.pass_obs_keys_to_actor})
         a_agent = self._sample_k_actions(a_dist, k=K_a)
         assert a_agent.shape == (K_a, B, L, G, D_action)
         if log_step:
@@ -742,6 +747,7 @@ def __init__(
         use_multigamma: bool = True,
         actor_type: Type[actor_critic.BaseActorHead] = actor_critic.Actor,
         critic_type: Type[actor_critic.BaseCriticHead] = actor_critic.NCriticsTwoHot,
+        pass_obs_keys_to_actor: Optional[Iterable[str]] = None,
     ):
         super().__init__(
             obs_space=obs_space,
@@ -766,6 +772,7 @@ def __init__(
             popart=popart,
             actor_type=actor_type,
             critic_type=critic_type,
+            pass_obs_keys_to_actor=pass_obs_keys_to_actor,
         )
 
     def _sample_k_actions(self, dist, k: int):
@@ -815,7 +822,7 @@ def forward(self, batch: Batch, log_step: bool):
         ################
         ## a ~ \pi(s) ##
         ################
-        a_dist = self.actor(s_rep, log_dict=active_log_dict)
+        a_dist = self.actor(s_rep, log_dict=active_log_dict, straight_from_obs={k : batch.obs[k] for k in self.pass_obs_keys_to_actor})
         if self.discrete:
             a_dist = DiscreteLikeContinuous(a_dist)
         if log_step:

amago/nets/actor_critic.py

@@ -49,7 +49,10 @@ def __init__(
         self.actions_differentiable = self.policy_dist.actions_differentiable
 
     def forward(
-        self, state: torch.Tensor, log_dict: Optional[dict] = None
+        self,
+        state: torch.Tensor,
+        log_dict: Optional[dict] = None,
+        straight_from_obs: Optional[dict[str, torch.Tensor]] = None,
     ) -> pyd.Distribution:
         """Compute an action distribution from a state representation.
 
@@ -61,7 +64,9 @@ def forward(
             (e.g. `Discrete` or `TanhGaussian`). Always a pytorch distribution (e.g., `Categorical`)
             where sampled actions would have shape (Batch, Length, Gammas, action_dim).
         """
-        dist_params = self.actor_network_forward(state=state, log_dict=log_dict)
+        dist_params = self.actor_network_forward(
+            state=state, log_dict=log_dict, straight_from_obs=straight_from_obs
+        )
         assert dist_params.ndim == 4
         assert dist_params.shape[-2:] == (
             self.num_gammas,
@@ -71,7 +76,10 @@ def forward(
 
     @abstractmethod
     def actor_network_forward(
-        self, state: torch.Tensor, log_dict: Optional[dict] = None
+        self,
+        state: torch.Tensor,
+        log_dict: Optional[dict] = None,
+        straight_from_obs: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
         raise NotImplementedError
 
@@ -127,7 +135,10 @@ def __init__(
         )
 
     def actor_network_forward(
-        self, state: torch.Tensor, log_dict: Optional[dict] = None
+        self,
+        state: torch.Tensor,
+        log_dict: Optional[dict] = None,
+        straight_from_obs: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
         dist_params = self.base(state)
         dist_params = rearrange(
@@ -216,7 +227,10 @@ def forward(self, x):
 
     @torch.compile
     def actor_network_forward(
-        self, state: torch.Tensor, log_dict: Optional[dict] = None
+        self,
+        state: torch.Tensor,
+        log_dict: Optional[dict] = None,
+        straight_from_obs: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
         B, L, D = state.shape
         x = self.inp(state)

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="amago",
-    version="3.1.0",
+    version="3.1.1",
     author="Jake Grigsby",
     author_email="grigsby@cs.utexas.edu",
     license="MIT",