Custom network integration

rosnav/model/agent_factory.py

@@ -54,7 +54,7 @@ def instantiate(cls, name: str, **kwargs) -> Union[Type[BaseAgent], Type[BasePol
         """
         assert name in cls.registry, f"Agent '{name}' is not registered!"
         agent_class = cls.registry[name]
-        
+
         if issubclass(agent_class, BaseAgent):
             return agent_class(**kwargs)
         else:

rosnav/model/custom_policy_from_json.py

@@ -0,0 +1,111 @@
+import os
+import rospy
+from typing import Callable, Dict, List, Optional, Tuple, Type, Union
+
+import gym
+import rospkg
+import torch as th
+import yaml
+import json
+
+from torch import nn
+from stable_baselines3.common.policies import ActorCriticPolicy
+
+
+from .agent_factory import AgentFactory
+from ..utils.utils import get_observation_space
+from .custom_policy_utils.utils import read_json
+from .custom_policy_utils.utils import create_body_network
+
+
+__all__ = ["CUSTOM"]
+
+
+""" 
+_RS: Robot state size - placeholder for robot related inputs to the NN
+_L: Number of laser beams - placeholder for the laser beam data 
+"""
+_L, _RS = get_observation_space()
+
+
+class CUSTOM_NETWORK(nn.Module):
+    """
+    Custom Multilayer Perceptron for policy and value function.
+
+    :param path: path to json file containing neural network
+    :param feature_dim: dimension of the features extracted with the features_extractor (e.g. features from a CNN)
+    :param last_layer_dim_pi: (int) number of units for the last layer of the policy network
+    :param last_layer_dim_vf: (int) number of units for the last layer of the value network
+    """
+
+    def __init__(
+        self,
+        path: str,
+        feature_dim: int,
+        last_layer_dim_pi: int = 32,
+        last_layer_dim_vf: int = 32,
+
+    ):
+        super(CUSTOM_NETWORK, self).__init__()
+
+        # Read file
+        data = read_json(path)
+
+        # Create the network based on JSON
+        self.body_net = create_body_network(data)
+
+        print(self.body_net)
+
+        # Save output dimensions, used to create the distributions
+        self.latent_dim_pi = last_layer_dim_pi
+        self.latent_dim_vf = last_layer_dim_vf
+
+        # Policy network
+        self.policy_net = nn.Sequential()
+
+        # Value network
+        self.value_net = nn.Sequential()
+
+    def forward(self, features: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network.
+            If all layers are shared, then ``latent_policy == latent_value``
+        """
+        body_x = self.body_net(features)
+        return self.policy_net(body_x), self.value_net(body_x)
+
+
+@AgentFactory.register("CUSTOM")
+class CUSTOM(ActorCriticPolicy):
+    """
+    Policy using the custom Multilayer Perceptron.
+    """
+
+    def __init__(
+        self,
+        observation_space: gym.spaces.Space,
+        action_space: gym.spaces.Space,
+        lr_schedule: Callable[[float], float],
+        net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None,
+        activation_fn: Type[nn.Module] = nn.ReLU,
+        *args,
+        **kwargs,
+
+    ):
+        # Getting the path to NN with rosparams
+        self.path=rospy.get_param("/custom_network_path")
+
+        super(CUSTOM, self).__init__(
+            observation_space,
+            action_space,
+            lr_schedule,
+            net_arch,
+            activation_fn,
+            *args,
+            **kwargs,
+        )
+        # Enable orthogonal initialization
+        self.ortho_init = True
+
+    def _build_mlp_extractor(self) -> None:
+        self.mlp_extractor = CUSTOM_NETWORK(self.path, 64)

rosnav/model/custom_policy_utils/modules.py

@@ -0,0 +1,28 @@
+from torch import nn
+
+def linear(params):
+    return nn.Linear(params["in_features"], params["out_features"], params.get("bias", True))
+
+def relu(params):
+    return nn.ReLU(params.get("inplace", False))
+
+def conv1d(params):
+    return nn.Conv1d(params["in_channels"], params["out_channels"], params["kernel_size"], params.get("stride", 1), params.get("padding", 0), params.get("dilation", 1), params.get("groups", 1), params.get("bias", True), params.get("padding_mode", 'zeros'))
+
+def conv2d(params):
+    return nn.Conv2d(params["in_channels"], params["out_channels"], params["kernel_size"], params.get("stride", 1), params.get("padding", 0), params.get("dilation", 1), params.get("groups", 1), params.get("bias", True), params.get("padding_mode", 'zeros'))
+
+def conv3d(params):
+    return nn.Conv3d(params["in_channels"], params["out_channels"], params["kernel_size"], params.get("stride", 1), params.get("padding", 0), params.get("dilation", 1), params.get("groups", 1), params.get("bias", True), params.get("padding_mode", 'zeros'))
+
+def tanh(params):
+    return nn.Tanh()
+
+modules = {
+    "relu" : relu,
+    "tanh" : tanh,
+    "linear" : linear,
+    "conv1d" : conv1d,
+    "conv2d" : conv2d,
+    "conv3d" : conv3d,
+}

rosnav/model/custom_policy_utils/utils.py

@@ -0,0 +1,24 @@
+import json
+from torch import nn
+from .modules import modules
+
+def read_json(path):
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def create_body_network(data):
+    body_net = nn.Sequential()
+
+    # Number of the module to be added to the NN
+    module_number = 0
+
+    # Iterate over each module given in the json in the policy
+    for module in data["policy"]:
+
+        # Add each module to the body network with corresponding number of module
+        body_net.add_module(f'{module_number}', modules[module["type"].lower()](module))
+
+        module_number += 1
+
+    return body_net