Skip to content

viberl.networks.base_network

Classes:

Name Description
BaseNetwork

Base neural network architecture for RL agents.

BaseNetwork

BaseNetwork(input_size: int, hidden_size: int = 128, num_hidden_layers: int = 2)

Bases: Module

Base neural network architecture for RL agents.

Methods:

Name Description
forward_backbone

Forward pass through the shared backbone.

init_weights

Initialize network weights using Xavier initialization.

Attributes:

Name Type Description
input_size
hidden_size
num_hidden_layers
backbone
Source code in viberl/networks/base_network.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def __init__(self, input_size: int, hidden_size: int = 128, num_hidden_layers: int = 2):
    super().__init__()
    self.input_size = input_size
    self.hidden_size = hidden_size
    self.num_hidden_layers = num_hidden_layers

    # Build network layers
    layers = []
    layers.append(nn.Linear(input_size, hidden_size))
    layers.append(nn.ReLU())

    for _ in range(num_hidden_layers - 1):
        layers.append(nn.Linear(hidden_size, hidden_size))
        layers.append(nn.ReLU())

    self.backbone = nn.Sequential(*layers)

input_size instance-attribute

input_size = input_size

hidden_size instance-attribute

hidden_size = hidden_size

num_hidden_layers instance-attribute

num_hidden_layers = num_hidden_layers

backbone instance-attribute

backbone = Sequential(*layers)

forward_backbone

forward_backbone(x: Tensor) -> Tensor

Forward pass through the shared backbone.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch_size, input_size)

required

Returns:

Type Description
Tensor

Processed tensor of shape (batch_size, hidden_size)

Source code in viberl/networks/base_network.py
25
26
27
28
29
30
31
32
33
34
def forward_backbone(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the shared backbone.

    Args:
        x: Input tensor of shape (batch_size, input_size)

    Returns:
        Processed tensor of shape (batch_size, hidden_size)
    """
    return self.backbone(x)

init_weights

init_weights() -> None

Initialize network weights using Xavier initialization.

Uses Xavier uniform initialization for linear layers and zeros for biases. This helps with stable gradient flow during training and prevents vanishing/exploding gradients.

Source code in viberl/networks/base_network.py
36
37
38
39
40
41
42
43
44
45
46
47
def init_weights(self) -> None:
    """Initialize network weights using Xavier initialization.

    Uses Xavier uniform initialization for linear layers and zeros for biases.
    This helps with stable gradient flow during training and prevents
    vanishing/exploding gradients.
    """
    for module in self.modules():
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.constant_(module.bias, 0)