Trpo tianshou

maro/cli/data_pipeline/citi_bike.py

@@ -8,7 +8,6 @@
 from enum import Enum
 
 import geopy.distance
-import numpy as np
 import pandas as pd
 from yaml import safe_load
 
@@ -320,7 +319,7 @@ def _process_distance(self, station_info: pd.DataFrame):
             0,
             index=station_info["station_index"],
             columns=station_info["station_index"],
-            dtype=np.float,
+            dtype=float,
         )
         look_up_df = station_info[["latitude", "longitude"]]
         return distance_adj.apply(
@@ -617,7 +616,7 @@ def _gen_distance(self, station_init: pd.DataFrame):
             0,
             index=station_init["station_index"],
             columns=station_init["station_index"],
-            dtype=np.float,
+            dtype=float,
         )
         look_up_df = station_init[["latitude", "longitude"]]
         distance_df = distance_adj.apply(

maro/rl/model/abs_net.py

@@ -47,6 +47,19 @@ def get_gradients(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
         loss.backward()
         return {name: param.grad for name, param in self.named_parameters()}
 
+    def get_kl_gradients(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Get the gradients with respect to all parameters according to the given loss.
+
+        Args:
+            loss (torch.tensor): Loss used to compute gradients.
+
+        Returns:
+            Gradients (Dict[str, torch.Tensor]): A dict that contains gradients for all parameters.
+        """
+        self.optim.zero_grad()
+        loss.backward(retain_graph=True)
+        return {name: param.grad for name, param in self.named_parameters()}
+
     def apply_gradients(self, grad: Dict[str, torch.Tensor]) -> None:
         """Apply gradients to the net to update all parameters.
 

maro/rl/policy/abs_policy.py

@@ -190,6 +190,18 @@ def get_gradients(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def get_kl_gradients(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Get the gradients with respect to all parameters of the internal nets according to the given loss.
+
+        Args:
+            loss (torch.tensor): Loss used to update the model.
+
+        Returns:
+            grad (Dict[str, torch.Tensor]): A dict that contains gradients of the internal nets for all parameters.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def apply_gradients(self, grad: dict) -> None:
         """Apply gradients to the net to update all parameters.

maro/rl/policy/discrete_rl_policy.py

@@ -286,6 +286,9 @@ def train_step(self, loss: torch.Tensor) -> None:
     def get_gradients(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
         return self._policy_net.get_gradients(loss)
 
+    def get_kl_gradients(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
+        return self._policy_net.get_kl_gradients(loss)
+
     def apply_gradients(self, grad: dict) -> None:
         self._policy_net.apply_gradients(grad)
 

maro/rl/training/algorithms/base/ac_ppo_base.py

@@ -309,3 +309,7 @@ async def train_step_as_task(self) -> None:
         for _ in range(self._params.grad_iters):
             if self._ops.update_actor_with_grad(await self._ops.get_actor_grad(batch)):  # early stop
                 break
+
+
+
+

maro/rl/training/algorithms/base/trpo_base.py

@@ -0,0 +1,45 @@
+import torch
+import torch.autograd as autograd
+import torch.nn as nn
+
+class Policy(nn.Module):
+    def __init__(self, num_inputs, num_outputs):
+        super(Policy, self).__init__()
+        self.affine1 = nn.Linear(num_inputs, 64)
+        self.affine2 = nn.Linear(64, 64)
+
+        self.action_mean = nn.Linear(64, num_outputs)
+        self.action_mean.weight.data.mul_(0.1)
+        self.action_mean.bias.data.mul_(0.0)
+
+        self.action_log_std = nn.Parameter(torch.zeros(1, num_outputs))
+
+        self.saved_actions = []
+        self.rewards = []
+        self.final_value = 0
+
+    def forward(self, x):
+        x = torch.tanh(self.affine1(x))
+        x = torch.tanh(self.affine2(x))
+
+        action_mean = self.action_mean(x)
+        action_log_std = self.action_log_std.expand_as(action_mean)
+        action_std = torch.exp(action_log_std)
+
+        return action_mean, action_log_std,action_std
+
+class Value(nn.Module):
+    def __init__(self, num_inputs):
+        super(Value, self).__init__()
+        self.affine1 = nn.Linear(num_inputs, 64)
+        self.affine2 = nn.Linear(64, 64)
+        self.value_head = nn.Linear(64, 1)
+        self.value_head.weight.data.mul_(0.1)
+        self.value_head.bias.data.mul_(0.0)
+
+    def forward(self, x):
+        x = torch.tanh(self.affine1(x))
+        x = torch.tanh(self.affine2(x))
+
+        state_values = self.value_head(x)
+        return state_values