Share more code between tf and torch policies

ml-agents/mlagents/trainers/policy/policy.py

 from abc import abstractmethod
+from typing import Dict, List, Optional
+import numpy as np
 
 from mlagents_envs.base_env import DecisionSteps
 from mlagents_envs.exception import UnityException
 
 
 class Policy(object):
-    def __init__(self, brain, seed):
+    def __init__(self, brain, seed, trainer_params):
+        self.use_continuous_act = brain.vector_action_space_type == "continuous"
+        if self.use_continuous_act:
+            self.num_branches = self.brain.vector_action_space_size[0]
+        else:
+            self.num_branches = len(self.brain.vector_action_space_size)
+        self.previous_action_dict: Dict[str, np.array] = {}
+        self.memory_dict: Dict[str, np.ndarray] = {}
+        self.normalize = trainer_params["normalize"]
+        self.use_recurrent = trainer_params["use_recurrent"]
+
+        if self.use_recurrent:
+            self.m_size = trainer_params["memory_size"]
+            self.sequence_length = trainer_params["sequence_length"]
+            if self.m_size == 0:
+                raise UnityPolicyException(
+                    "The memory size for brain {0} is 0 even "
+                    "though the trainer uses recurrent.".format(brain.brain_name)
+                )
+            elif self.m_size % 2 != 0:
+                raise UnityPolicyException(
+                    "The memory size for brain {0} is {1} "
+                    "but it must be divisible by 2.".format(
+                        brain.brain_name, self.m_size
+                    )
+                )
+
+    def make_empty_memory(self, num_agents):
+        """
+        Creates empty memory for use with RNNs
+        :param num_agents: Number of agents.
+        :return: Numpy array of zeros.
+        """
+        return np.zeros((num_agents, self.m_size), dtype=np.float32)
+
+    def save_memories(
+        self, agent_ids: List[str], memory_matrix: Optional[np.ndarray]
+    ) -> None:
+        if memory_matrix is None:
+            return
+        for index, agent_id in enumerate(agent_ids):
+            self.memory_dict[agent_id] = memory_matrix[index, :]
+
+    def retrieve_memories(self, agent_ids: List[str]) -> np.ndarray:
+        memory_matrix = np.zeros((len(agent_ids), self.m_size), dtype=np.float32)
+        for index, agent_id in enumerate(agent_ids):
+            if agent_id in self.memory_dict:
+                memory_matrix[index, :] = self.memory_dict[agent_id]
+        return memory_matrix
+
+    def remove_memories(self, agent_ids):
+        for agent_id in agent_ids:
+            if agent_id in self.memory_dict:
+                self.memory_dict.pop(agent_id)
+
+    def make_empty_previous_action(self, num_agents):
+        """
+        Creates empty previous action for use with RNNs and discrete control
+        :param num_agents: Number of agents.
+        :return: Numpy array of zeros.
+        """
+        return np.zeros((num_agents, self.num_branches), dtype=np.int)
+
+    def save_previous_action(
+        self, agent_ids: List[str], action_matrix: Optional[np.ndarray]
+    ) -> None:
+        if action_matrix is None:
+            return
+        for index, agent_id in enumerate(agent_ids):
+            self.previous_action_dict[agent_id] = action_matrix[index, :]
+
+    def retrieve_previous_action(self, agent_ids: List[str]) -> np.ndarray:
+        action_matrix = np.zeros((len(agent_ids), self.num_branches), dtype=np.int)
+        for index, agent_id in enumerate(agent_ids):
+            if agent_id in self.previous_action_dict:
+                action_matrix[index, :] = self.previous_action_dict[agent_id]
+        return action_matrix
+
+    def remove_previous_action(self, agent_ids):
+        for agent_id in agent_ids:
+            if agent_id in self.previous_action_dict:
+                self.previous_action_dict.pop(agent_id)
 
     def get_action(
         self, decision_requests: DecisionSteps, worker_id: int = 0

ml-agents/mlagents/trainers/policy/tf_policy.py

         self.vec_obs_size = brain.vector_observation_space_size
         self.vis_obs_size = brain.number_visual_observations
 
-        self.use_recurrent = trainer_parameters["use_recurrent"]
-        self.memory_dict: Dict[str, np.ndarray] = {}
-        self.num_branches = len(self.brain.vector_action_space_size)
-        self.previous_action_dict: Dict[str, np.array] = {}
-        self.normalize = trainer_parameters.get("normalize", False)
-        self.use_continuous_act = brain.vector_action_space_type == "continuous"
-        if self.use_continuous_act:
-            self.num_branches = self.brain.vector_action_space_size[0]
         self.model_path = trainer_parameters["model_path"]
         self.initialize_path = trainer_parameters.get("init_path", None)
         self.keep_checkpoints = trainer_parameters.get("keep_checkpoints", 5)
         )
         self.saver = None
         self.seed = seed
-        if self.use_recurrent:
-            self.m_size = trainer_parameters["memory_size"]
-            self.sequence_length = trainer_parameters["sequence_length"]
-            if self.m_size == 0:
-                raise UnityPolicyException(
-                    "The memory size for brain {0} is 0 even "
-                    "though the trainer uses recurrent.".format(brain.brain_name)
-                )
-            elif self.m_size % 2 != 0:
-                raise UnityPolicyException(
-                    "The memory size for brain {0} is {1} "
-                    "but it must be divisible by 2.".format(
-                        brain.brain_name, self.m_size
-                    )
-                )
         self._initialize_tensorflow_references()
         self.load = load
 
                 mask = 1 - np.concatenate(batched_step_result.action_mask, axis=1)
             feed_dict[self.action_masks] = mask
         return feed_dict
-
-    def make_empty_memory(self, num_agents):
-        """
-        Creates empty memory for use with RNNs
-        :param num_agents: Number of agents.
-        :return: Numpy array of zeros.
-        """
-        return np.zeros((num_agents, self.m_size), dtype=np.float32)
-
-    def save_memories(
-        self, agent_ids: List[str], memory_matrix: Optional[np.ndarray]
-    ) -> None:
-        if memory_matrix is None:
-            return
-        for index, agent_id in enumerate(agent_ids):
-            self.memory_dict[agent_id] = memory_matrix[index, :]
-
-    def retrieve_memories(self, agent_ids: List[str]) -> np.ndarray:
-        memory_matrix = np.zeros((len(agent_ids), self.m_size), dtype=np.float32)
-        for index, agent_id in enumerate(agent_ids):
-            if agent_id in self.memory_dict:
-                memory_matrix[index, :] = self.memory_dict[agent_id]
-        return memory_matrix
-
-    def remove_memories(self, agent_ids):
-        for agent_id in agent_ids:
-            if agent_id in self.memory_dict:
-                self.memory_dict.pop(agent_id)
-
-    def make_empty_previous_action(self, num_agents):
-        """
-        Creates empty previous action for use with RNNs and discrete control
-        :param num_agents: Number of agents.
-        :return: Numpy array of zeros.
-        """
-        return np.zeros((num_agents, self.num_branches), dtype=np.int)
-
-    def save_previous_action(
-        self, agent_ids: List[str], action_matrix: Optional[np.ndarray]
-    ) -> None:
-        if action_matrix is None:
-            return
-        for index, agent_id in enumerate(agent_ids):
-            self.previous_action_dict[agent_id] = action_matrix[index, :]
-
-    def retrieve_previous_action(self, agent_ids: List[str]) -> np.ndarray:
-        action_matrix = np.zeros((len(agent_ids), self.num_branches), dtype=np.int)
-        for index, agent_id in enumerate(agent_ids):
-            if agent_id in self.previous_action_dict:
-                action_matrix[index, :] = self.previous_action_dict[agent_id]
-        return action_matrix
-
-    def remove_previous_action(self, agent_ids):
-        for agent_id in agent_ids:
-            if agent_id in self.previous_action_dict:
-                self.previous_action_dict.pop(agent_id)
 
     def get_current_step(self):
         """

ml-agents/mlagents/trainers/policy/torch_policy.py

         :param reparameterize: Whether we are using the resampling trick to update the policy
         in continuous output.
         """
-        super(TorchPolicy, self).__init__(brain, seed)
+        super(TorchPolicy, self).__init__(brain, seed, trainer_params)
-        self.normalize = trainer_params["normalize"]
         self.seed = seed
         self.brain = brain
         self.global_step = 0
         If this policy normalizes vector observations, this will update the norm values in the graph.
         :param vector_obs: The vector observations to add to the running estimate of the distribution.
         """
+        print(vector_obs.shape)
+        vector_obs = [vector_obs]
-            self.critic.network_body.normalize(vector_obs)
-            self.actor.network_body.normalize(vector_obs)
+            self.critic.network_body.update_normalization(vector_obs)
+            self.actor.network_body.update_normalization(vector_obs)
 
     def execute_model(self, vec_obs, vis_obs, masks=None):
         action_dists = self.actor(vec_obs, vis_obs, masks)
             actions.append(action)
             log_probs.append(action_dist.log_prob(action))
             entropies.append(action_dist.entropy())
-        actions = torch.stack(actions)
-        log_probs = torch.stack(log_probs)
-        entropies = torch.stack(entropies)
+        actions = torch.stack(actions).squeeze(0)
+        log_probs = torch.stack(log_probs).squeeze(0)
+        entropies = torch.stack(entropies).squeeze(0)
 
         value_heads, mean_value = self.critic(vec_obs, vis_obs)
         return actions, log_probs, entropies, value_heads
         :return: Outputs from network as defined by self.inference_dict.
         """
         vec_obs, vis_obs, masks = self.split_decision_step(decision_requests)
-        vec_obs = [vec_obs]  # For consistency with visual observations
+        vec_obs = [torch.Tensor(vec_obs)]
+        vis_obs = [torch.Tensor(vis_ob) for vis_ob in vis_obs]
+        run_out["pre_action"] = np.array(
+            action.detach()
+        )  # Todo - make pre_action difference
         run_out["log_probs"] = np.array(log_probs.detach())
         run_out["entropy"] = np.array(entropy.detach())
         run_out["value_heads"] = {
     @property
     def use_vec_obs(self):
         return self.vec_obs_size > 0
-
-    @property
-    def use_recurrent(self):
-        return False
-
-    @property
-    def use_continuous_act(self):
-        return True
 
     def get_current_step(self):
         """