update(hus): update data from official server

ppof_ch4_code_p1.py

@@ -0,0 +1,324 @@
+# pip install minigrid
+from typing import Union, Tuple, Dict, List, Optional
+from multiprocessing import Process
+import multiprocessing as mp
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import minigrid
+import gymnasium as gym
+from torch.optim.lr_scheduler import ExponentialLR, MultiStepLR
+from tensorboardX import SummaryWriter
+from minigrid.wrappers import FlatObsWrapper
+
+random.seed(0)
+np.random.seed(0)
+torch.manual_seed(0)
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+else:
+    device = torch.device("cpu")
+
+train_config = dict(
+    train_iter=1024,
+    train_data_count=128,
+    test_data_count=4096,
+)
+
+little_RND_net_config = dict(
+    exp_name="little_rnd_network",
+    observation_shape=2835,
+    hidden_size_list=[32, 16],
+    learning_rate=1e-3,
+    batch_size=64,
+    update_per_collect=100,
+    obs_norm=True,
+    obs_norm_clamp_min=-1,
+    obs_norm_clamp_max=1,
+    reward_mse_ratio=1e5,
+)
+
+small_RND_net_config = dict(
+    exp_name="small_rnd_network",
+    observation_shape=2835,
+    hidden_size_list=[64, 64],
+    learning_rate=1e-3,
+    batch_size=64,
+    update_per_collect=100,
+    obs_norm=True,
+    obs_norm_clamp_min=-1,
+    obs_norm_clamp_max=1,
+    reward_mse_ratio=1e5,
+)
+
+standard_RND_net_config = dict(
+    exp_name="standard_rnd_network",
+    observation_shape=2835,
+    hidden_size_list=[128, 64],
+    learning_rate=1e-3,
+    batch_size=64,
+    update_per_collect=100,
+    obs_norm=True,
+    obs_norm_clamp_min=-1,
+    obs_norm_clamp_max=1,
+    reward_mse_ratio=1e5,
+)
+
+large_RND_net_config = dict(
+    exp_name="large_RND_network",
+    observation_shape=2835,
+    hidden_size_list=[256, 256],
+    learning_rate=1e-3,
+    batch_size=64,
+    update_per_collect=100,
+    obs_norm=True,
+    obs_norm_clamp_min=-1,
+    obs_norm_clamp_max=1,
+    reward_mse_ratio=1e5,
+)
+
+very_large_RND_net_config = dict(
+    exp_name="very_large_RND_network",
+    observation_shape=2835,
+    hidden_size_list=[512, 512],
+    learning_rate=1e-3,
+    batch_size=64,
+    update_per_collect=100,
+    obs_norm=True,
+    obs_norm_clamp_min=-1,
+    obs_norm_clamp_max=1,
+    reward_mse_ratio=1e5,
+)
+
+class FCEncoder(nn.Module):
+    def __init__(
+            self,
+            obs_shape: int,
+            hidden_size_list,
+            activation: Optional[nn.Module] = nn.ReLU(),
+    ) -> None:
+        super(FCEncoder, self).__init__()
+        self.obs_shape = obs_shape
+        self.act = activation
+        self.init = nn.Linear(obs_shape, hidden_size_list[0])
+
+        layers = []
+        for i in range(len(hidden_size_list) - 1):
+            layers.append(nn.Linear(hidden_size_list[i], hidden_size_list[i + 1]))
+            layers.append(self.act)
+        self.main = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act(self.init(x))
+        x = self.main(x)
+        return x
+
+class RndNetwork(nn.Module):
+    def __init__(self, obs_shape: Union[int, list], hidden_size_list: list) -> None:
+        super(RndNetwork, self).__init__()
+        self.target = FCEncoder(obs_shape, hidden_size_list)
+        self.predictor = FCEncoder(obs_shape, hidden_size_list)
+
+        for param in self.target.parameters():
+            param.requires_grad = False
+
+    def forward(self, obs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        predict_feature = self.predictor(obs)
+        with torch.no_grad():
+            target_feature = self.target(obs)
+        return predict_feature, target_feature
+
+class RunningMeanStd(object):
+    def __init__(self, epsilon=1e-4, shape=(), device=torch.device('cpu')):
+        self._epsilon = epsilon
+        self._shape = shape
+        self._device = device
+        self.reset()
+
+    def update(self, x):
+        batch_mean = np.mean(x, axis=0)
+        batch_var = np.var(x, axis=0)
+        batch_count = x.shape[0]
+
+        new_count = batch_count + self._count
+        mean_delta = batch_mean - self._mean
+        new_mean = self._mean + mean_delta * batch_count / new_count
+        # this method for calculating new variable might be numerically unstable
+        m_a = self._var * self._count
+        m_b = batch_var * batch_count
+        m2 = m_a + m_b + np.square(mean_delta) * self._count * batch_count / new_count
+        new_var = m2 / new_count
+        self._mean = new_mean
+        self._var = new_var
+        self._count = new_count
+
+    def reset(self):
+        if len(self._shape) > 0:
+            self._mean = np.zeros(self._shape, 'float32')
+            self._var = np.ones(self._shape, 'float32')
+        else:
+            self._mean, self._var = 0., 1.
+        self._count = self._epsilon
+
+    @property
+    def mean(self) -> np.ndarray:
+        if np.isscalar(self._mean):
+            return self._mean
+        else:
+            return torch.FloatTensor(self._mean).to(self._device)
+
+    @property
+    def std(self) -> np.ndarray:
+        std = np.sqrt(self._var + 1e-8)
+        if np.isscalar(std):
+            return std
+        else:
+            return torch.FloatTensor(std).to(self._device)
+
+class RndRewardModel():
+
+    def __init__(self, config) -> None:  # noqa
+        super(RndRewardModel, self).__init__()
+        self.cfg = config
+
+        self.tb_logger = SummaryWriter(config["exp_name"])
+        self.reward_model = RndNetwork(
+            obs_shape=config["observation_shape"], hidden_size_list=config["hidden_size_list"]
+        ).to(device)
+
+        self.opt = optim.Adam(self.reward_model.predictor.parameters(), config["learning_rate"])
+        self.scheduler = ExponentialLR(self.opt, gamma=0.997)
+
+        self.estimate_cnt_rnd = 0
+        if self.cfg["obs_norm"]:
+            self._running_mean_std_rnd_obs = RunningMeanStd(epsilon=1e-4, device=device)
+
+    def __del__(self):
+        self.tb_logger.flush()
+        self.tb_logger.close()
+
+    def train(self, data) -> None:
+        for _ in range(self.cfg["update_per_collect"]):
+            train_data: list = random.sample(data, self.cfg["batch_size"])
+            train_data: torch.Tensor = torch.stack(train_data).to(device)
+            if self.cfg["obs_norm"]:
+                # Note: observation normalization: transform obs to mean 0, std 1
+                self._running_mean_std_rnd_obs.update(train_data.cpu().numpy())
+                train_data = (train_data - self._running_mean_std_rnd_obs.mean) / self._running_mean_std_rnd_obs.std
+                train_data = torch.clamp(
+                    train_data, min=self.cfg["obs_norm_clamp_min"], max=self.cfg["obs_norm_clamp_max"]
+                )
+
+            predict_feature, target_feature = self.reward_model(train_data)
+            loss = F.mse_loss(predict_feature, target_feature.detach())
+            self.opt.zero_grad()
+            loss.backward()
+            self.opt.step()
+        self.scheduler.step()
+
+    def estimate(self, data: list) -> List[Dict]:
+        """
+        estimate the rnd intrinsic reward
+        """
+
+        obs = torch.stack(data).to(device)
+        if self.cfg["obs_norm"]:
+            # Note: observation normalization: transform obs to mean 0, std 1
+            obs = (obs - self._running_mean_std_rnd_obs.mean) / self._running_mean_std_rnd_obs.std
+            obs = torch.clamp(obs, min=self.cfg["obs_norm_clamp_min"], max=self.cfg["obs_norm_clamp_max"])
+
+        with torch.no_grad():
+            self.estimate_cnt_rnd += 1
+            predict_feature, target_feature = self.reward_model(obs)
+            mse = F.mse_loss(predict_feature, target_feature, reduction='none').mean(dim=1)
+            self.tb_logger.add_scalar('rnd_reward/mse', mse.cpu().numpy().mean(), self.estimate_cnt_rnd)
+
+            # Note: according to the min-max normalization, transform rnd reward to [0,1]
+            rnd_reward = mse * self.cfg["reward_mse_ratio"]  #(mse - mse.min()) / (mse.max() - mse.min() + 1e-11)
+
+            self.tb_logger.add_scalar('rnd_reward/rnd_reward_max', rnd_reward.max(), self.estimate_cnt_rnd)
+            self.tb_logger.add_scalar('rnd_reward/rnd_reward_mean', rnd_reward.mean(), self.estimate_cnt_rnd)
+            self.tb_logger.add_scalar('rnd_reward/rnd_reward_min', rnd_reward.min(), self.estimate_cnt_rnd)
+
+            rnd_reward = torch.chunk(rnd_reward, rnd_reward.shape[0], dim=0)
+
+def training(config, train_data, test_data):
+    rnd_reward_model = RndRewardModel(config=config)
+    for i in range(train_config["train_iter"]):
+        rnd_reward_model.train([torch.Tensor(item["last_observation"]) for item in train_data[i]])
+        rnd_reward_model.estimate([torch.Tensor(item["last_observation"]) for item in test_data])
+
+def main():
+    env = gym.make("MiniGrid-Empty-8x8-v0")
+    env_obs = FlatObsWrapper(env)
+
+    train_data = []
+    test_data = []
+
+    for i in range(train_config["train_iter"]):
+
+        train_data_per_iter = []
+
+        while len(train_data_per_iter) < train_config["train_data_count"]:
+            last_observation, _ = env_obs.reset()
+            terminated = False
+            while terminated != True and len(train_data_per_iter) < train_config["train_data_count"]:
+                action = env_obs.action_space.sample()
+                observation, reward, terminated, truncated, info = env_obs.step(action)
+                train_data_per_iter.append(
+                    {
+                        "last_observation": last_observation,
+                        "action": action,
+                        "reward": reward,
+                        "observation": observation
+                    }
+                )
+                last_observation = observation
+            env_obs.close()
+
+        train_data.append(train_data_per_iter)
+
+    while len(test_data) < train_config["test_data_count"]:
+        last_observation, _ = env_obs.reset()
+        terminated = False
+        while terminated != True and len(train_data_per_iter) < train_config["test_data_count"]:
+            action = env_obs.action_space.sample()
+            observation, reward, terminated, truncated, info = env_obs.step(action)
+            test_data.append(
+                {
+                    "last_observation": last_observation,
+                    "action": action,
+                    "reward": reward,
+                    "observation": observation
+                }
+            )
+            last_observation = observation
+        env_obs.close()
+
+    p0 = Process(target=training, args=(little_RND_net_config, train_data, test_data))
+    p0.start()
+
+    p1 = Process(target=training, args=(small_RND_net_config, train_data, test_data))
+    p1.start()
+
+    p2 = Process(target=training, args=(standard_RND_net_config, train_data, test_data))
+    p2.start()
+
+    p3 = Process(target=training, args=(large_RND_net_config, train_data, test_data))
+    p3.start()
+
+    p4 = Process(target=training, args=(very_large_RND_net_config, train_data, test_data))
+    p4.start()
+
+    p0.join()
+    p1.join()
+    p2.join()
+    p3.join()
+    p4.join()
+
+if __name__ == "__main__":
+    mp.set_start_method('spawn')
+    main()

ppof_ch4_data_lunarlander.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff98aa71827552cd72afc108edddac8e1d77df3499c624dc6f16e256b2a79d61
+size 99443

ppof_ch4_data_p1.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c993afb3adb533830ae271f86ba9fb587e70216385f6f20e88dab7fa8f583d8
+size 4035833

ppof_ch5_code_p1.py

@@ -0,0 +1,193 @@
+"""
+Long Short Term Memory (LSTM) <link https://ieeexplore.ieee.org/abstract/document/6795963 link> is a kind of recurrent neural network that can capture long-short term information.
+This document mainly includes:
+- Pytorch implementation for LSTM.
+- An example to test LSTM.
+For beginners, you can refer to <link https://zhuanlan.zhihu.com/p/32085405 link> to learn the basics about how LSTM works.
+"""
+from typing import Optional, Union, Tuple, List, Dict
+import math
+import torch
+import torch.nn as nn
+from ding.torch_utils import build_normalization
+
+
+class LSTM(nn.Module):
+    """
+    **Overview:**
+        Implementation of LSTM cell with layer norm.
+    """
+
+    def __init__(
+            self,
+            input_size: int,
+            hidden_size: int,
+            num_layers: int,
+            norm_type: Optional[str] = 'LN',
+            dropout: float = 0.
+    ) -> None:
+        # Initialize arguments.
+        super(LSTM, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        # Initialize normalization functions.
+        norm_func = build_normalization(norm_type)
+        self.norm = nn.ModuleList([norm_func(hidden_size * 4) for _ in range(2 * num_layers)])
+        # Initialize LSTM parameters.
+        self.wx = nn.ParameterList()
+        self.wh = nn.ParameterList()
+        dims = [input_size] + [hidden_size] * num_layers
+        for l in range(num_layers):
+            self.wx.append(nn.Parameter(torch.zeros(dims[l], dims[l + 1] * 4)))
+            self.wh.append(nn.Parameter(torch.zeros(hidden_size, hidden_size * 4)))
+        self.bias = nn.Parameter(torch.zeros(num_layers, hidden_size * 4))
+        # Initialize the Dropout Layer.
+        self.use_dropout = dropout > 0.
+        if self.use_dropout:
+            self.dropout = nn.Dropout(dropout)
+        self._init()
+
+    # Dealing with different types of input and return preprocessed prev_state.
+    def _before_forward(self, inputs: torch.Tensor, prev_state: Union[None, List[Dict]]) -> torch.Tensor:
+        seq_len, batch_size = inputs.shape[:2]
+        # If prev_state is None, it indicates that this is the beginning of a sequence. In this case, prev_state will be initialized as zero.
+        if prev_state is None:
+            zeros = torch.zeros(self.num_layers, batch_size, self.hidden_size, dtype=inputs.dtype, device=inputs.device)
+            prev_state = (zeros, zeros)
+        # If prev_state is not None, then preprocess it into one batch.
+        else:
+            assert len(prev_state) == batch_size
+            state = [[v for v in prev.values()] for prev in prev_state]
+            state = list(zip(*state))
+            prev_state = [torch.cat(t, dim=1) for t in state]
+
+        return prev_state
+
+    def _init(self):
+        # Initialize parameters. Each parameter is initialized using a uniform distribution of: $$U(-\sqrt {\frac 1 {HiddenSize}}, -\sqrt {\frac 1 {HiddenSize}})$$
+        gain = math.sqrt(1. / self.hidden_size)
+        for l in range(self.num_layers):
+            torch.nn.init.uniform_(self.wx[l], -gain, gain)
+            torch.nn.init.uniform_(self.wh[l], -gain, gain)
+            if self.bias is not None:
+                torch.nn.init.uniform_(self.bias[l], -gain, gain)
+
+    def forward(
+            self,
+            inputs: torch.Tensor,
+            prev_state: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, list]]:
+        # The shape of input is: [sequence length, batch size, input size]
+        seq_len, batch_size = inputs.shape[:2]
+        prev_state = self._before_forward(inputs, prev_state)
+
+        H, C = prev_state
+        x = inputs
+        next_state = []
+        for l in range(self.num_layers):
+            h, c = H[l], C[l]
+            new_x = []
+            for s in range(seq_len):
+                # Calculate $$z, z^i, z^f, z^o$$ simultaneously.
+                gate = self.norm[l * 2](torch.matmul(x[s], self.wx[l])
+                                        ) + self.norm[l * 2 + 1](torch.matmul(h, self.wh[l]))
+                if self.bias is not None:
+                    gate += self.bias[l]
+                gate = list(torch.chunk(gate, 4, dim=1))
+                i, f, o, z = gate
+                # $$z^i = \sigma (Wx^ix^t + Wh^ih^{t-1})$$
+                i = torch.sigmoid(i)
+                # $$z^f = \sigma (Wx^fx^t + Wh^fh^{t-1})$$
+                f = torch.sigmoid(f)
+                # $$z^o = \sigma (Wx^ox^t + Wh^oh^{t-1})$$
+                o = torch.sigmoid(o)
+                # $$z = tanh(Wxx^t + Whh^{t-1})$$
+                z = torch.tanh(z)
+                # $$c^t = z^f \odot c^{t-1}+z^i \odot z$$
+                c = f * c + i * z
+                # $$h^t = z^o \odot tanh(c^t)$$
+                h = o * torch.tanh(c)
+                new_x.append(h)
+            next_state.append((h, c))
+            x = torch.stack(new_x, dim=0)
+            # Dropout layer.
+            if self.use_dropout and l != self.num_layers - 1:
+                x = self.dropout(x)
+        next_state = [torch.stack(t, dim=0) for t in zip(*next_state)]
+        # Return list type, split the next_state .
+        h, c = next_state
+        batch_size = h.shape[1]
+        # Split h with shape [num_layers, batch_size, hidden_size] to a list with length batch_size and each element is a tensor with shape [num_layers, 1, hidden_size]. The same operation is performed on c.
+        next_state = [torch.chunk(h, batch_size, dim=1), torch.chunk(c, batch_size, dim=1)]
+        next_state = list(zip(*next_state))
+        next_state = [{k: v for k, v in zip(['h', 'c'], item)} for item in next_state]
+        return x, next_state
+
+
+def pack_data(data: List[torch.Tensor], traj_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Overview:
+        You need to pack variable-length data to regular tensor, return tensor and corresponding mask.
+        If len(data_i) < traj_len, use `null_padding`,
+        else split the whole sequences info different trajectories.
+    Returns:
+        - tensor (:obj:`torch.Tensor`): dtype (torch.float32), shape (traj_len, B, N)
+        - mask (:obj:`torch.Tensor`): dtype (torch.float32), shape (traj_len, B)
+    """
+    new_data = []
+    mask = []
+    for item in data:
+        D, N = item.shape
+        if D < traj_len:
+            null_padding = torch.zeros(traj_len - D, N)
+            new_item = torch.cat([item, null_padding])
+            new_data.append(new_item)
+            item_mask = torch.ones(traj_len)
+            item_mask[D:].zero_()
+            mask.append(item_mask)
+        else:
+            for i in range(0, D, traj_len):
+                item_mask = torch.ones(traj_len)
+                new_item = item[i:i + traj_len]
+                if new_item.shape[0] < traj_len:
+                    new_item = item[-traj_len:]
+                new_data.append(new_item)
+                mask.append(torch.ones(traj_len))
+    new_data = torch.stack(new_data, dim=1)
+    mask = torch.stack(mask, dim=1)
+
+    return new_data, mask
+
+
+def test_lstm():
+    seq_len_list = [32, 49, 24, 78, 45]
+    traj_len = 32
+    N = 10
+    hidden_size = 32
+    num_layers = 2
+
+    variable_len_data = [torch.rand(s, N) for s in seq_len_list]
+    input_, mask = pack_data(variable_len_data, traj_len)
+    assert isinstance(input_, torch.Tensor), type(input_)
+    batch_size = input_.shape[1]
+    assert batch_size == 9, "packed data must have 9 trajectories"
+    lstm = LSTM(N, hidden_size=hidden_size, num_layers=num_layers, norm_type='LN', dropout=0.1)
+
+    prev_state = None
+    for s in range(traj_len):
+        input_step = input_[s:s + 1]
+        output, prev_state = lstm(input_step, prev_state)
+
+    assert output.shape == (1, batch_size, hidden_size)
+    assert len(prev_state) == batch_size
+    assert prev_state[0]['h'].shape == (num_layers, 1, hidden_size)
+    loss = (output * mask.unsqueeze(-1)).mean()
+    loss.backward()
+    for _, m in lstm.named_parameters():
+        assert isinstance(m.grad, torch.Tensor)
+    print('finished')
+
+
+if __name__ == '__main__':
+    test_lstm()

ppof_ch6_code_p1.py

@@ -0,0 +1,79 @@
+import numpy as np
+import torch
+
+
+def get_agent_id_feature(agent_id, agent_num):
+    agent_id_feature = torch.zeros(agent_num)
+    agent_id_feature[agent_id] = 1
+    return agent_id_feature
+
+
+def get_movement_feature():
+    # for simplicity, we use random movement feature here
+    movement_feature = torch.randint(0, 2, (8, ))
+    return movement_feature
+
+
+def get_own_feature():
+    # for simplicity, we use random own feature here
+    return torch.randn(10)
+
+
+def get_ally_visible_feature():
+    # this function only return the visible feature of one ally
+    # for simplicity, we use random tensor as ally visible feature while zero tensor as ally invisible feature
+    if np.random.random() > 0.5:
+        ally_visible_feature = torch.randn(4)
+    else:
+        ally_visible_feature = torch.zeros(4)
+    return ally_visible_feature
+
+
+def get_enemy_visible_feature():
+    # this function only return the visible feature of one enemy
+    # for simplicity, we use random tensor as enemy visible feature while zero tensor as enemy invisible feature
+    if np.random.random() > 0.8:
+        enemy_visible_feature = torch.randn(4)
+    else:
+        enemy_visible_feature = torch.zeros(4)
+    return enemy_visible_feature
+
+
+def get_ind_global_state(agent_id, ally_agent_num, enemy_agent_num):
+    # You need to implement this function
+    raise NotImplementedError
+
+
+def get_ep_global_state(agent_id, ally_agent_num, enemy_agent_num):
+    # In many multi-agent environments such as SMAC, the global state is the simplified version of the combination
+    # of all the agent's independent state, and the concrete implementation depends on the characteris of environment.
+    # For simplicity, we use random feature here.
+    ally_center_feature = torch.randn(8)
+    enemy_center_feature = torch.randn(8)
+    return torch.cat([ally_center_feature, enemy_center_feature])
+
+
+def get_as_global_state(agent_id, ally_agent_num, enemy_agent_num):
+    # You need to implement this function
+    raise NotImplementedError
+
+
+def test_global_state():
+    ally_agent_num = 3
+    enemy_agent_num = 5
+    # get independent global state, which usually used in decentralized training
+    for agent_id in range(ally_agent_num):
+        ind_global_state = get_ind_global_state(agent_id, ally_agent_num, enemy_agent_num)
+        assert isinstance(ind_global_state, torch.Tensor)
+    # get environment provide global state, which is the same for all agents, used in centralized training
+    for agent_id in range(ally_agent_num):
+        ep_global_state = get_ep_global_state(agent_id, ally_agent_num, enemy_agent_num)
+        assert isinstance(ep_global_state, torch.Tensor)
+    # get naive agent-specific global state, which is the specific for each agent, used in centralized training
+    for agent_id in range(ally_agent_num):
+        as_global_state = get_as_global_state(agent_id, ally_agent_num, enemy_agent_num)
+        assert isinstance(as_global_state, torch.Tensor)
+
+
+if __name__ == "__main__":
+    test_global_state()

ppof_ch7_code_p1.py

@@ -0,0 +1,114 @@
+from typing import Optional, Tuple, List
+import torch
+import torch.nn as nn
+import treetensor.torch as ttorch
+
+
+class PPOFModel(nn.Module):
+    mode = ['compute_actor', 'compute_critic', 'compute_actor_critic']
+
+    def __init__(
+            self,
+            obs_shape: Tuple[int],
+            action_shape: int,
+            encoder_hidden_size_list: List = [128, 128, 64],
+            actor_head_hidden_size: int = 64,
+            actor_head_layer_num: int = 1,
+            critic_head_hidden_size: int = 64,
+            critic_head_layer_num: int = 1,
+            activation: Optional[nn.Module] = nn.ReLU(),
+    ) -> None:
+        super(PPOFModel, self).__init__()
+        self.obs_shape, self.action_shape = obs_shape, action_shape
+
+        # encoder
+        layers = []
+        input_size = obs_shape[0]
+        kernel_size_list = [8, 4, 3]
+        stride_list = [4, 2, 1]
+        for i in range(len(encoder_hidden_size_list)):
+            output_size = encoder_hidden_size_list[i]
+            layers.append(nn.Conv2d(input_size, output_size, kernel_size_list[i], stride_list[i]))
+            layers.append(activation)
+            input_size = output_size
+        layers.append(nn.Flatten())
+        self.encoder = nn.Sequential(*layers)
+
+        flatten_size = input_size = self.get_flatten_size()
+        # critic
+        layers = []
+        for i in range(critic_head_layer_num):
+            layers.append(nn.Linear(input_size, critic_head_hidden_size))
+            layers.append(activation)
+            input_size = critic_head_hidden_size
+        layers.append(nn.Linear(critic_head_hidden_size, 1))
+        self.critic = nn.Sequential(*layers)
+        # actor
+        layers = []
+        input_size = flatten_size
+        for i in range(actor_head_layer_num):
+            layers.append(nn.Linear(input_size, actor_head_hidden_size))
+            layers.append(activation)
+            input_size = actor_head_hidden_size
+        self.actor = nn.Sequential(*layers)
+        self.mu = nn.Linear(actor_head_hidden_size, action_shape)
+        self.log_sigma = nn.Parameter(torch.zeros(1, action_shape))
+
+        # init weights
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        # You need to implement this function
+        raise NotImplementedError
+
+    def get_flatten_size(self) -> int:
+        test_data = torch.randn(1, *self.obs_shape)
+        with torch.no_grad():
+            output = self.encoder(test_data)
+        return output.shape[1]
+
+    def forward(self, inputs: ttorch.Tensor, mode: str) -> ttorch.Tensor:
+        assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
+        return getattr(self, mode)(inputs)
+
+    def compute_actor(self, x: ttorch.Tensor) -> ttorch.Tensor:
+        x = self.encoder(x)
+        x = self.actor(x)
+        mu = self.mu(x)
+        log_sigma = self.log_sigma + torch.zeros_like(mu)  # addition aims to broadcast shape
+        sigma = torch.exp(log_sigma)
+        return ttorch.as_tensor({'mu': mu, 'sigma': sigma})
+
+    def compute_critic(self, x: ttorch.Tensor) -> ttorch.Tensor:
+        x = self.encoder(x)
+        value = self.critic(x)
+        return value
+
+    def compute_actor_critic(self, x: ttorch.Tensor) -> ttorch.Tensor:
+        x = self.encoder(x)
+        value = self.critic(x)
+        x = self.actor(x)
+        mu = self.mu(x)
+        log_sigma = self.log_sigma + torch.zeros_like(mu)  # addition aims to broadcast shape
+        sigma = torch.exp(log_sigma)
+        return ttorch.as_tensor({'logit': {'mu': mu, 'sigma': sigma}, 'value': value})
+
+
+def test_ppof_model() -> None:
+    model = PPOFModel((4, 84, 84), 5)
+    print(model)
+    data = torch.randn(3, 4, 84, 84)
+    output = model(data, mode='compute_critic')
+    assert output.shape == (3, 1)
+    output = model(data, mode='compute_actor')
+    assert output.mu.shape == (3, 5)
+    assert output.sigma.shape == (3, 5)
+    output = model(data, mode='compute_actor_critic')
+    assert output.value.shape == (3, 1)
+    assert output.logit.mu.shape == (3, 5)
+    assert output.logit.sigma.shape == (3, 5)
+    print('End...')
+
+
+if __name__ == "__main__":
+    test_ppof_model()