Neural Network Modules¶

Critic Network¶

CriticNetwork ¶

CriticNetwork(
    encoder: Module,
    value_head: Module | None = None,
    embed_dim: int = 128,
    hidden_dim: int = 512,
    customized: bool = False,
)

Bases: Module

Create a critic network given an encoder (e.g. as the one in the policy network) with a value head to transform the embeddings to a scalar value.

Parameters:

encoder (Module) –

Encoder module to encode the input
value_head (Module | None, default: None ) –

Value head to transform the embeddings to a scalar value
embed_dim (int, default: 128 ) –

Dimension of the embeddings of the value head
hidden_dim (int, default: 512 ) –

Dimension of the hidden layer of the value head

Methods:

forward –

Forward pass of the critic network: encode the imput in embedding space and return the value

Source code in rl4co/models/rl/common/critic.py

def __init__(
    self,
    encoder: nn.Module,
    value_head: nn.Module | None = None,
    embed_dim: int = 128,
    hidden_dim: int = 512,
    customized: bool = False,
):
    super().__init__()

    self.encoder = encoder
    if value_head is None:
        # check if embed dim of encoder is different, if so, use it
        if getattr(encoder, "embed_dim", embed_dim) != embed_dim:
            log.warning(
                f"Found encoder with different embed_dim {encoder.embed_dim} than the value head {embed_dim}. \
                Using encoder embed_dim for value head."
            )
            embed_dim = getattr(encoder, "embed_dim", embed_dim)
        value_head = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1)
        )
    self.value_head = value_head
    self.customized = customized

forward ¶

forward(x: Tensor | TensorDict, hidden=None) -> Tensor

Forward pass of the critic network: encode the imput in embedding space and return the value

Parameters:

x (Tensor | TensorDict) –

Input containing the environment state. Can be a Tensor or a TensorDict

Returns:

Tensor –

Value of the input state

Source code in rl4co/models/rl/common/critic.py

def forward(self, x: Tensor | TensorDict, hidden=None) -> Tensor:
    """Forward pass of the critic network: encode the imput in embedding space and return the value

    Args:
        x: Input containing the environment state. Can be a Tensor or a TensorDict

    Returns:
        Value of the input state
    """
    if not self.customized:  # fir for most of costructive tasks
        h, _ = self.encoder(x)  # [batch_size, N, embed_dim] -> [batch_size, N]
        return self.value_head(h).mean(1)  # [batch_size, N] -> [batch_size]
    else:  # custimized encoder and value head with hidden input
        h = self.encoder(x)  # [batch_size, N, embed_dim] -> [batch_size, N]
        return self.value_head(h, hidden)

Graph Neural Networks¶

MultiHeadAttentionLayer ¶

MultiHeadAttentionLayer(
    embed_dim: int,
    num_heads: int = 8,
    feedforward_hidden: int = 512,
    normalization: str | None = "batch",
    bias: bool = True,
    sdpa_fn: Callable | None = None,
    moe_kwargs: dict | None = None,
)

Bases: Sequential

Multi-Head Attention Layer with normalization and feed-forward layer

Parameters:

embed_dim (int) –

dimension of the embeddings
num_heads (int, default: 8 ) –

number of heads in the MHA
feedforward_hidden (int, default: 512 ) –

dimension of the hidden layer in the feed-forward layer
normalization (str | None, default: 'batch' ) –

type of normalization to use (batch, layer, none)
sdpa_fn (Callable | None, default: None ) –

scaled dot product attention function (SDPA)
moe_kwargs (dict | None, default: None ) –

Keyword arguments for MoE

Source code in rl4co/models/nn/graph/attnnet.py

def __init__(
    self,
    embed_dim: int,
    num_heads: int = 8,
    feedforward_hidden: int = 512,
    normalization: str | None = "batch",
    bias: bool = True,
    sdpa_fn: Callable | None = None,
    moe_kwargs: dict | None = None,
):
    num_neurons = [feedforward_hidden] if feedforward_hidden > 0 else []
    if moe_kwargs is not None:
        ffn = MoE(embed_dim, embed_dim, num_neurons=num_neurons, **moe_kwargs)
    else:
        ffn = MLP(
            input_dim=embed_dim,
            output_dim=embed_dim,
            num_neurons=num_neurons,
            hidden_act="ReLU",
        )

    super().__init__(
        SkipConnection(MultiHeadAttention(embed_dim, num_heads, bias=bias, sdpa_fn=sdpa_fn)),
        Normalization(embed_dim, normalization),
        SkipConnection(ffn),
        Normalization(embed_dim, normalization),
    )

GraphAttentionNetwork ¶

GraphAttentionNetwork(
    num_heads: int,
    embed_dim: int,
    num_layers: int,
    normalization: str = "batch",
    feedforward_hidden: int = 512,
    sdpa_fn: Callable | None = None,
    moe_kwargs: dict | None = None,
)

Bases: Module

Graph Attention Network to encode embeddings with a series of MHA layers consisting of a MHA layer, normalization, feed-forward layer, and normalization. Similar to Transformer encoder, as used in Kool et al. (2019).

Parameters:

num_heads (int) –

number of heads in the MHA
embed_dim (int) –

dimension of the embeddings
num_layers (int) –

number of MHA layers
normalization (str, default: 'batch' ) –

type of normalization to use (batch, layer, none)
feedforward_hidden (int, default: 512 ) –

dimension of the hidden layer in the feed-forward layer
sdpa_fn (Callable | None, default: None ) –

scaled dot product attention function (SDPA)
moe_kwargs (dict | None, default: None ) –

Keyword arguments for MoE

Methods:

forward –

Forward pass of the encoder

Source code in rl4co/models/nn/graph/attnnet.py

def __init__(
    self,
    num_heads: int,
    embed_dim: int,
    num_layers: int,
    normalization: str = "batch",
    feedforward_hidden: int = 512,
    sdpa_fn: Callable | None = None,
    moe_kwargs: dict | None = None,
):
    super().__init__()

    self.layers = nn.Sequential(
        *(
            MultiHeadAttentionLayer(
                embed_dim,
                num_heads,
                feedforward_hidden=feedforward_hidden,
                normalization=normalization,
                sdpa_fn=sdpa_fn,
                moe_kwargs=moe_kwargs,
            )
            for _ in range(num_layers)
        )
    )

forward ¶

forward(x: Tensor, mask: Tensor | None = None) -> Tensor

Forward pass of the encoder

Parameters:

x (Tensor) –

[batch_size, graph_size, embed_dim] initial embeddings to process
mask (Tensor | None, default: None ) –

[batch_size, graph_size, graph_size] mask for the input embeddings. Unused for now.

Source code in rl4co/models/nn/graph/attnnet.py

def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
    """Forward pass of the encoder

    Args:
        x: [batch_size, graph_size, embed_dim] initial embeddings to process
        mask: [batch_size, graph_size, graph_size] mask for the input embeddings. Unused for now.
    """
    assert mask is None, "Mask not yet supported!"
    h = self.layers(x)
    return h

GCNEncoder ¶

GCNEncoder(
    env_name: str,
    embed_dim: int,
    num_layers: int,
    init_embedding: Module = None,
    residual: bool = True,
    edge_idx_fn: EdgeIndexFnSignature = None,
    dropout: float = 0.5,
    bias: bool = True,
)

Bases: Module

Graph Convolutional Network to encode embeddings with a series of GCN layers from the pytorch geometric package

Parameters:

embed_dim (int) –

dimension of the embeddings
num_nodes –

number of nodes in the graph
num_gcn_layer –

number of GCN layers
self_loop –

whether to add self loop in the graph
residual (bool, default: True ) –

whether to use residual connection

Methods:

forward –

Forward pass of the encoder.

Source code in rl4co/models/nn/graph/gcn.py

def __init__(
    self,
    env_name: str,
    embed_dim: int,
    num_layers: int,
    init_embedding: nn.Module = None,
    residual: bool = True,
    edge_idx_fn: EdgeIndexFnSignature = None,
    dropout: float = 0.5,
    bias: bool = True,
):
    super().__init__()

    self.env_name = env_name
    self.embed_dim = embed_dim
    self.residual = residual
    self.dropout = dropout

    self.init_embedding = (
        env_init_embedding(self.env_name, {"embed_dim": embed_dim})
        if init_embedding is None
        else init_embedding
    )

    if edge_idx_fn is None:
        log.warning("No edge indices passed. Assume a fully connected graph")
        edge_idx_fn = edge_idx_fn_wrapper

    self.edge_idx_fn = edge_idx_fn

    # Define the GCN layers
    self.gcn_layers = nn.ModuleList(
        [GCNConv(embed_dim, embed_dim, bias=bias) for _ in range(num_layers)]
    )

forward ¶

forward(
    td: TensorDict, mask: Tensor | None = None
) -> tuple[Tensor, Tensor]

Forward pass of the encoder. Transform the input TensorDict into a latent representation.

Parameters:

td (TensorDict) –

Input TensorDict containing the environment state
mask (Tensor | None, default: None ) –

Mask to apply to the attention

Returns:

h ( Tensor ) –

Latent representation of the input
init_h ( Tensor ) –

Initial embedding of the input

Source code in rl4co/models/nn/graph/gcn.py

def forward(self, td: TensorDict, mask: Tensor | None = None) -> tuple[Tensor, Tensor]:
    """Forward pass of the encoder.
    Transform the input TensorDict into a latent representation.

    Args:
        td: Input TensorDict containing the environment state
        mask: Mask to apply to the attention

    Returns:
        h: Latent representation of the input
        init_h: Initial embedding of the input
    """
    # Transfer to embedding space
    init_h = self.init_embedding(td)
    bs, num_nodes, emb_dim = init_h.shape
    # (bs*num_nodes, emb_dim)
    update_node_feature = init_h.reshape(-1, emb_dim)
    # shape=(2, num_edges)
    edge_index = self.edge_idx_fn(td, num_nodes)

    for layer in self.gcn_layers[:-1]:
        update_node_feature = layer(update_node_feature, edge_index)
        update_node_feature = F.relu(update_node_feature)
        update_node_feature = F.dropout(
            update_node_feature, training=self.training, p=self.dropout
        )

    # last layer without relu activation and dropout
    update_node_feature = self.gcn_layers[-1](update_node_feature, edge_index)

    # De-batch the graph
    update_node_feature = update_node_feature.view(bs, num_nodes, emb_dim)

    # Residual
    if self.residual:
        update_node_feature = update_node_feature + init_h

    return update_node_feature, init_h

MessagePassingEncoder ¶

MessagePassingEncoder(
    env_name: str,
    embed_dim: int,
    num_nodes: int,
    num_layers: int,
    init_embedding: Module = None,
    aggregation: str = "add",
    self_loop: bool = False,
    residual: bool = True,
)

Bases: Module

Source code in rl4co/models/nn/graph/mpnn.py

def __init__(
    self,
    env_name: str,
    embed_dim: int,
    num_nodes: int,
    num_layers: int,
    init_embedding: nn.Module = None,
    aggregation: str = "add",
    self_loop: bool = False,
    residual: bool = True,
):
    """
    Note:
        - Support fully connected graph for now.
    """
    super().__init__()

    self.env_name = env_name

    self.init_embedding = (
        env_init_embedding(self.env_name, {"embed_dim": embed_dim})
        if init_embedding is None
        else init_embedding
    )

    # Generate edge index for a fully connected graph
    adj_matrix = torch.ones(num_nodes, num_nodes)
    if self_loop:
        adj_matrix.fill_diagonal_(0)  # No self-loops
    self.edge_index = torch.permute(torch.nonzero(adj_matrix), (1, 0))

    # Init message passing models
    self.mpnn_layers = nn.ModuleList(
        [
            MessagePassingLayer(
                node_indim=embed_dim,
                node_outdim=embed_dim,
                edge_indim=1,
                edge_outdim=1,
                aggregation=aggregation,
                residual=residual,
            )
            for _ in range(num_layers)
        ]
    )

    # Record parameters
    self.self_loop = self_loop

Attention Mechanisms¶

Classes:

MultiHeadAttention –

PyTorch native implementation of Flash Multi-Head Attention with automatic mixed precision support.
MultiHeadCrossAttention –

PyTorch native implementation of Flash Multi-Head Cross Attention with automatic mixed precision support.
PointerAttention –

Calculate logits given query, key and value and logit key.
PointerAttnMoE –

Calculate logits given query, key and value and logit key.
MultiHeadCompat –
PolyNetAttention –

Calculate logits given query, key and value and logit key.

Functions:

scaled_dot_product_attention_simple –

Simple (exact) Scaled Dot-Product Attention in RL4CO without customized kernels (i.e. no Flash Attention).

MultiHeadAttention ¶

MultiHeadAttention(
    embed_dim: int,
    num_heads: int,
    bias: bool = True,
    attention_dropout: float = 0.0,
    causal: bool = False,
    device: str = None,
    dtype: dtype = None,
    sdpa_fn: Callable | None = None,
)

Bases: Module

PyTorch native implementation of Flash Multi-Head Attention with automatic mixed precision support. Uses PyTorch's native scaled_dot_product_attention implementation, available from 2.0

Note

If scaled_dot_product_attention is not available, use custom implementation of scaled_dot_product_attention without Flash Attention.

Parameters:

embed_dim (int) –

total dimension of the model
num_heads (int) –

number of heads
bias (bool, default: True ) –

whether to use bias
attention_dropout (float, default: 0.0 ) –

dropout rate for attention weights
causal (bool, default: False ) –

whether to apply causal mask to attention scores
device (str, default: None ) –

torch device
dtype (dtype, default: None ) –

torch dtype
sdpa_fn (Callable | None, default: None ) –

scaled dot product attention function (SDPA) implementation

Methods:

forward –

x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim)

Source code in rl4co/models/nn/attention.py

def __init__(
    self,
    embed_dim: int,
    num_heads: int,
    bias: bool = True,
    attention_dropout: float = 0.0,
    causal: bool = False,
    device: str = None,
    dtype: torch.dtype = None,
    sdpa_fn: Callable | None = None,
) -> None:
    factory_kwargs = {"device": device, "dtype": dtype}
    super().__init__()
    self.embed_dim = embed_dim
    self.causal = causal
    self.attention_dropout = attention_dropout
    self.sdpa_fn = sdpa_fn if sdpa_fn is not None else scaled_dot_product_attention

    self.num_heads = num_heads
    assert self.embed_dim % num_heads == 0, "self.kdim must be divisible by num_heads"
    self.head_dim = self.embed_dim // num_heads
    assert self.head_dim % 8 == 0 and self.head_dim <= 128, (
        "Only support head_dim <= 128 and divisible by 8"
    )

    self.Wqkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs)
    self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)

forward ¶

forward(x, attn_mask=None)

x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) attn_mask: bool tensor of shape (batch, seqlen)

Source code in rl4co/models/nn/attention.py

def forward(self, x, attn_mask=None):
    """x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim)
    attn_mask: bool tensor of shape (batch, seqlen)
    """
    # Project query, key, value
    q, k, v = rearrange(
        self.Wqkv(x), "b s (three h d) -> three b h s d", three=3, h=self.num_heads
    ).unbind(dim=0)

    if attn_mask is not None:
        attn_mask = (
            attn_mask.unsqueeze(1)
            if attn_mask.ndim == 3
            else attn_mask.unsqueeze(1).unsqueeze(2)
        )

    # Scaled dot product attention
    out = self.sdpa_fn(
        q,
        k,
        v,
        attn_mask=attn_mask,
        dropout_p=self.attention_dropout,
    )
    return self.out_proj(rearrange(out, "b h s d -> b s (h d)"))

MultiHeadCrossAttention ¶

MultiHeadCrossAttention(
    embed_dim: int,
    num_heads: int,
    bias: bool = False,
    attention_dropout: float = 0.0,
    device: str = None,
    dtype: dtype = None,
    sdpa_fn: Callable | Module | None = None,
)

Bases: Module

PyTorch native implementation of Flash Multi-Head Cross Attention with automatic mixed precision support. Uses PyTorch's native scaled_dot_product_attention implementation, available from 2.0

Note

If scaled_dot_product_attention is not available, use custom implementation of scaled_dot_product_attention without Flash Attention.

Parameters:

embed_dim (int) –

total dimension of the model
num_heads (int) –

number of heads
bias (bool, default: False ) –

whether to use bias
attention_dropout (float, default: 0.0 ) –

dropout rate for attention weights
device (str, default: None ) –

torch device
dtype (dtype, default: None ) –

torch dtype
sdpa_fn (Callable | Module | None, default: None ) –

scaled dot product attention function (SDPA)

Source code in rl4co/models/nn/attention.py

def __init__(
    self,
    embed_dim: int,
    num_heads: int,
    bias: bool = False,
    attention_dropout: float = 0.0,
    device: str = None,
    dtype: torch.dtype = None,
    sdpa_fn: Callable | nn.Module | None = None,
) -> None:
    factory_kwargs = {"device": device, "dtype": dtype}
    super().__init__()
    self.embed_dim = embed_dim
    self.attention_dropout = attention_dropout

    # Default to `scaled_dot_product_attention` if `sdpa_fn` is not provided
    if sdpa_fn is None:
        sdpa_fn = sdpa_fn_wrapper
    self.sdpa_fn = sdpa_fn

    self.num_heads = num_heads
    assert self.embed_dim % num_heads == 0, "self.kdim must be divisible by num_heads"
    self.head_dim = self.embed_dim // num_heads
    assert self.head_dim % 8 == 0 and self.head_dim <= 128, (
        "Only support head_dim <= 128 and divisible by 8"
    )

    self.Wq = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)
    self.Wkv = nn.Linear(embed_dim, 2 * embed_dim, bias=bias, **factory_kwargs)
    self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)

PointerAttention ¶

PointerAttention(
    embed_dim: int,
    num_heads: int,
    mask_inner: bool = True,
    out_bias: bool = False,
    check_nan: bool = True,
    sdpa_fn: Callable | str = "default",
    **kwargs
)

Bases: Module

Calculate logits given query, key and value and logit key. This follows the pointer mechanism of Vinyals et al. (2015) (https://arxiv.org/abs/1506.03134).

Note

With Flash Attention, masking is not supported

Performs the following

Apply cross attention to get the heads
Project heads to get glimpse
Compute attention score between glimpse and logit key

Parameters:

embed_dim (int) –

total dimension of the model
num_heads (int) –

number of heads
mask_inner (bool, default: True ) –

whether to mask inner attention
linear_bias –

whether to use bias in linear projection
check_nan (bool, default: True ) –

whether to check for NaNs in logits
sdpa_fn (Callable | str, default: 'default' ) –

scaled dot product attention function (SDPA) implementation

Methods:

forward –

Compute attention logits given query, key, value, logit key and attention mask.

Source code in rl4co/models/nn/attention.py

def __init__(
    self,
    embed_dim: int,
    num_heads: int,
    mask_inner: bool = True,
    out_bias: bool = False,
    check_nan: bool = True,
    sdpa_fn: Callable | str = "default",
    **kwargs,
):
    super().__init__()
    self.num_heads = num_heads
    self.mask_inner = mask_inner

    # Projection - query, key, value already include projections
    self.project_out = nn.Linear(embed_dim, embed_dim, bias=out_bias)
    self.check_nan = check_nan

    # Defaults for sdpa_fn implementation
    # see https://github.com/ai4co/rl4co/issues/228
    if isinstance(sdpa_fn, str):
        if sdpa_fn == "default":
            sdpa_fn = scaled_dot_product_attention
        elif sdpa_fn == "simple":
            sdpa_fn = scaled_dot_product_attention_simple
        else:
            raise ValueError(
                f"Unknown sdpa_fn: {sdpa_fn}. Available options: ['default', 'simple']"
            )
    else:
        if sdpa_fn is None:
            sdpa_fn = scaled_dot_product_attention
            log.info("Using default scaled_dot_product_attention for PointerAttention")
    self.sdpa_fn = sdpa_fn

forward ¶

forward(query, key, value, logit_key, attn_mask=None)

Compute attention logits given query, key, value, logit key and attention mask.

Parameters:

query –

query tensor of shape [B, ..., L, E]
key –

key tensor of shape [B, ..., S, E]
value –

value tensor of shape [B, ..., S, E]
logit_key –

logit key tensor of shape [B, ..., S, E]
attn_mask –

attention mask tensor of shape [B, ..., S]. Note that True means that the value should take part in attention as described in the PyTorch Documentation

Source code in rl4co/models/nn/attention.py

def forward(self, query, key, value, logit_key, attn_mask=None):
    """Compute attention logits given query, key, value, logit key and attention mask.

    Args:
        query: query tensor of shape [B, ..., L, E]
        key: key tensor of shape [B, ..., S, E]
        value: value tensor of shape [B, ..., S, E]
        logit_key: logit key tensor of shape [B, ..., S, E]
        attn_mask: attention mask tensor of shape [B, ..., S]. Note that `True` means that the value _should_ take part in attention
            as described in the [PyTorch Documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
    """
    # Compute inner multi-head attention with no projections.
    heads = self._inner_mha(query, key, value, attn_mask)
    glimpse = self._project_out(heads, attn_mask)

    # Batch matrix multiplication to compute logits (batch_size, num_steps, graph_size)
    # bmm is slightly faster than einsum and matmul
    logits = (torch.bmm(glimpse, logit_key.squeeze(-2).transpose(-2, -1))).squeeze(
        -2
    ) / math.sqrt(glimpse.size(-1))

    if self.check_nan:
        assert not torch.isnan(logits).any(), "Logits contain NaNs"

    return logits

PointerAttnMoE ¶

PointerAttnMoE(
    embed_dim: int,
    num_heads: int,
    mask_inner: bool = True,
    out_bias: bool = False,
    check_nan: bool = True,
    sdpa_fn: Callable | None = None,
    moe_kwargs: dict | None = None,
)

Bases: PointerAttention

Calculate logits given query, key and value and logit key. This follows the pointer mechanism of Vinyals et al. (2015) https://arxiv.org/abs/1506.03134, and the MoE gating mechanism of Zhou et al. (2024) https://arxiv.org/abs/2405.01029.

Note

With Flash Attention, masking is not supported

Performs the following

Apply cross attention to get the heads
Project heads to get glimpse
Compute attention score between glimpse and logit key

Parameters:

embed_dim (int) –

total dimension of the model
num_heads (int) –

number of heads
mask_inner (bool, default: True ) –

whether to mask inner attention
linear_bias –

whether to use bias in linear projection
check_nan (bool, default: True ) –

whether to check for NaNs in logits
sdpa_fn (Callable | None, default: None ) –

scaled dot product attention function (SDPA) implementation
moe_kwargs (dict | None, default: None ) –

Keyword arguments for MoE

Source code in rl4co/models/nn/attention.py

def __init__(
    self,
    embed_dim: int,
    num_heads: int,
    mask_inner: bool = True,
    out_bias: bool = False,
    check_nan: bool = True,
    sdpa_fn: Callable | None = None,
    moe_kwargs: dict | None = None,
):
    super().__init__(embed_dim, num_heads, mask_inner, out_bias, check_nan, sdpa_fn)
    self.moe_kwargs = moe_kwargs

    self.project_out = None
    self.project_out_moe = MoE(
        embed_dim, embed_dim, num_neurons=[], out_bias=out_bias, **moe_kwargs
    )
    if self.moe_kwargs["light_version"]:
        self.dense_or_moe = nn.Linear(embed_dim, 2, bias=False)
        self.project_out = nn.Linear(embed_dim, embed_dim, bias=out_bias)

MultiHeadCompat ¶

MultiHeadCompat(
    n_heads,
    input_dim,
    embed_dim=None,
    val_dim=None,
    key_dim=None,
)

Bases: Module

Methods:

forward –

:param q: queries (batch_size, n_query, input_dim)

Source code in rl4co/models/nn/attention.py

def __init__(self, n_heads, input_dim, embed_dim=None, val_dim=None, key_dim=None):
    super().__init__()

    if val_dim is None:
        # assert embed_dim is not None, "Provide either embed_dim or val_dim"
        val_dim = embed_dim // n_heads
    if key_dim is None:
        key_dim = val_dim

    self.n_heads = n_heads
    self.input_dim = input_dim
    self.embed_dim = embed_dim
    self.val_dim = val_dim
    self.key_dim = key_dim

    self.W_query = nn.Parameter(torch.Tensor(n_heads, input_dim, key_dim))
    self.W_key = nn.Parameter(torch.Tensor(n_heads, input_dim, key_dim))

    self.init_parameters()

forward ¶

forward(q, h=None, mask=None)

:param q: queries (batch_size, n_query, input_dim) :param h: data (batch_size, graph_size, input_dim) :param mask: mask (batch_size, n_query, graph_size) or viewable as that (i.e. can be 2 dim if n_query == 1) Mask should contain 1 if attention is not possible (i.e. mask is negative adjacency) :return:

Source code in rl4co/models/nn/attention.py

def forward(self, q, h=None, mask=None):
    """

    :param q: queries (batch_size, n_query, input_dim)
    :param h: data (batch_size, graph_size, input_dim)
    :param mask: mask (batch_size, n_query, graph_size) or viewable as that (i.e. can be 2 dim if n_query == 1)
    Mask should contain 1 if attention is not possible (i.e. mask is negative adjacency)
    :return:
    """

    if h is None:
        h = q  # compute self-attention

    # h should be (batch_size, graph_size, input_dim)
    batch_size, graph_size, input_dim = h.size()
    n_query = q.size(1)

    hflat = h.contiguous().view(-1, input_dim)  #################   reshape
    qflat = q.contiguous().view(-1, input_dim)

    # last dimension can be different for keys and values
    shp = (self.n_heads, batch_size, graph_size, -1)
    shp_q = (self.n_heads, batch_size, n_query, -1)

    # Calculate queries, (n_heads, n_query, graph_size, key/val_size)
    Q = torch.matmul(qflat, self.W_query).view(shp_q)
    K = torch.matmul(hflat, self.W_key).view(shp)

    # Calculate compatibility (n_heads, batch_size, n_query, graph_size)
    compatibility_s2n = torch.matmul(Q, K.transpose(2, 3))

    return compatibility_s2n

PolyNetAttention ¶

PolyNetAttention(
    k: int,
    embed_dim: int,
    poly_layer_dim: int,
    num_heads: int,
    **kwargs
)

Bases: PointerAttention

Calculate logits given query, key and value and logit key. This implements a modified version the pointer mechanism of Vinyals et al. (2015) (https://arxiv.org/abs/1506.03134) as described in Hottung et al. (2024) (https://arxiv.org/abs/2402.14048) PolyNetAttention conditions the attention logits on a set of k different binary vectors allowing to learn k different solution strategies.

Note

With Flash Attention, masking is not supported

Performs the following

Apply cross attention to get the heads
Project heads to get glimpse
Apply PolyNet layers
Compute attention score between glimpse and logit key

Parameters:

k (int) –

Number unique bit vectors used to compute attention score
embed_dim (int) –

total dimension of the model
poly_layer_dim (int) –

Dimension of the PolyNet layers
num_heads (int) –

number of heads
mask_inner –

whether to mask inner attention
linear_bias –

whether to use bias in linear projection
check_nan –

whether to check for NaNs in logits
sdpa_fn –

scaled dot product attention function (SDPA) implementation

Methods:

forward –

Compute attention logits given query, key, value, logit key and attention mask.

Source code in rl4co/models/nn/attention.py

def __init__(self, k: int, embed_dim: int, poly_layer_dim: int, num_heads: int, **kwargs):
    super().__init__(embed_dim, num_heads, **kwargs)

    self.k = k
    self.binary_vector_dim = math.ceil(math.log2(k))
    self.binary_vectors = torch.nn.Parameter(
        torch.Tensor(list(itertools.product([0, 1], repeat=self.binary_vector_dim))[:k]),
        requires_grad=False,
    )

    self.poly_layer_1 = nn.Linear(embed_dim + self.binary_vector_dim, poly_layer_dim)
    self.poly_layer_2 = nn.Linear(poly_layer_dim, embed_dim)

forward ¶

forward(query, key, value, logit_key, attn_mask=None)

Compute attention logits given query, key, value, logit key and attention mask.

Parameters:

query –

query tensor of shape [B, ..., L, E]
key –

key tensor of shape [B, ..., S, E]
value –

value tensor of shape [B, ..., S, E]
logit_key –

logit key tensor of shape [B, ..., S, E]
attn_mask –

attention mask tensor of shape [B, ..., S]. Note that True means that the value should take part in attention as described in the PyTorch Documentation

Source code in rl4co/models/nn/attention.py

def forward(self, query, key, value, logit_key, attn_mask=None):
    """Compute attention logits given query, key, value, logit key and attention mask.

    Args:
        query: query tensor of shape [B, ..., L, E]
        key: key tensor of shape [B, ..., S, E]
        value: value tensor of shape [B, ..., S, E]
        logit_key: logit key tensor of shape [B, ..., S, E]
        attn_mask: attention mask tensor of shape [B, ..., S]. Note that `True` means that the value _should_ take part in attention
            as described in the [PyTorch Documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
    """
    # Compute inner multi-head attention with no projections.
    heads = self._inner_mha(query, key, value, attn_mask)
    glimpse = self.project_out(heads)

    num_solutions = glimpse.shape[1]
    z = self.binary_vectors.repeat(math.ceil(num_solutions / self.k), 1)[:num_solutions]
    z = z[None].expand(glimpse.shape[0], num_solutions, self.binary_vector_dim)

    # PolyNet layers
    poly_out = self.poly_layer_1(torch.cat((glimpse, z), dim=2))
    poly_out = F.relu(poly_out)
    poly_out = self.poly_layer_2(poly_out)

    glimpse += poly_out

    # Batch matrix multiplication to compute logits (batch_size, num_steps, graph_size)
    # bmm is slightly faster than einsum and matmul
    logits = (torch.bmm(glimpse, logit_key.squeeze(-2).transpose(-2, -1))).squeeze(
        -2
    ) / math.sqrt(glimpse.size(-1))

    if self.check_nan:
        assert not torch.isnan(logits).any(), "Logits contain NaNs"

    return logits

scaled_dot_product_attention_simple ¶

scaled_dot_product_attention_simple(
    q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False
)

Simple (exact) Scaled Dot-Product Attention in RL4CO without customized kernels (i.e. no Flash Attention).

Source code in rl4co/models/nn/attention.py

def scaled_dot_product_attention_simple(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False):
    """Simple (exact) Scaled Dot-Product Attention in RL4CO without customized kernels (i.e. no Flash Attention)."""

    # Check for causal and attn_mask conflict
    if is_causal and attn_mask is not None:
        raise ValueError("Cannot set both is_causal and attn_mask")

    # Calculate scaled dot product
    scores = torch.matmul(q, k.transpose(-2, -1)) / (k.size(-1) ** 0.5)

    # Apply the provided attention mask
    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            scores.masked_fill_(~attn_mask, float("-inf"))
        else:
            scores += attn_mask

    # Apply causal mask
    if is_causal:
        s, l_ = scores.size(-2), scores.size(-1)
        mask = torch.triu(torch.ones((s, l_), device=scores.device), diagonal=1)
        scores.masked_fill_(mask.bool(), float("-inf"))

    # Softmax to get attention weights
    attn_weights = F.softmax(scores, dim=-1)

    # Apply dropout
    if dropout_p > 0.0:
        attn_weights = F.dropout(attn_weights, p=dropout_p)

    # Compute the weighted sum of values
    return torch.matmul(attn_weights, v)

Multi-Layer Perceptron¶

MLP ¶

MLP(
    input_dim: int,
    output_dim: int,
    num_neurons: list[int] = [64, 32],
    dropout_probs: None | list[float] = None,
    hidden_act: str = "ReLU",
    out_act: str = "Identity",
    input_norm: str = "None",
    output_norm: str = "None",
)

Bases: Module

Source code in rl4co/models/nn/mlp.py

def __init__(
    self,
    input_dim: int,
    output_dim: int,
    num_neurons: list[int] = [64, 32],
    dropout_probs: None | list[float] = None,
    hidden_act: str = "ReLU",
    out_act: str = "Identity",
    input_norm: str = "None",
    output_norm: str = "None",
):
    super().__init__()

    assert input_norm in ["Batch", "Layer", "None"]
    assert output_norm in ["Batch", "Layer", "None"]

    if dropout_probs is None:
        dropout_probs = [0.0] * len(num_neurons)
    elif len(dropout_probs) != len(num_neurons):
        log.info(
            "dropout_probs List length should match the num_neurons List length for MLP, dropouts set to False instead"
        )
        dropout_probs = [0.0] * len(num_neurons)

    self.input_dim = input_dim
    self.output_dim = output_dim
    self.num_neurons = num_neurons
    self.hidden_act = getattr(nn, hidden_act)()
    self.out_act = getattr(nn, out_act)()
    self.dropouts = []
    for i in range(len(dropout_probs)):
        self.dropouts.append(nn.Dropout(p=dropout_probs[i]))

    input_dims = [input_dim] + num_neurons
    output_dims = num_neurons + [output_dim]

    self.lins = nn.ModuleList()
    for i, (in_dim, out_dim) in enumerate(zip(input_dims, output_dims)):
        self.lins.append(nn.Linear(in_dim, out_dim))

    self.input_norm = self._get_norm_layer(input_norm, input_dim)
    self.output_norm = self._get_norm_layer(output_norm, output_dim)

Operations¶

PositionalEncoding ¶

PositionalEncoding(
    embed_dim: int,
    dropout: float = 0.1,
    max_len: int = 1000,
)

Bases: Module

Methods:

forward –

Arguments:

Source code in rl4co/models/nn/ops.py

def __init__(self, embed_dim: int, dropout: float = 0.1, max_len: int = 1000):
    super().__init__()
    self.dropout = nn.Dropout(p=dropout)
    self.d_model = embed_dim
    max_len = max_len
    position = torch.arange(max_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, self.d_model, 2) * (-math.log(10000.0) / self.d_model))
    pe = torch.zeros(max_len, 1, self.d_model)
    pe[:, 0, 0::2] = torch.sin(position * div_term)
    pe[:, 0, 1::2] = torch.cos(position * div_term)
    pe = pe.transpose(0, 1)  # [1, max_len, d_model]
    self.register_buffer("pe", pe)

forward ¶

forward(hidden: Tensor, seq_pos) -> Tensor

Parameters:

x –

Tensor, shape [batch_size, seq_len, embedding_dim]
seq_pos –

Tensor, shape [batch_size, seq_len]

Source code in rl4co/models/nn/ops.py

def forward(self, hidden: torch.Tensor, seq_pos) -> torch.Tensor:
    """
    Arguments:
        x: Tensor, shape ``[batch_size, seq_len, embedding_dim]``
        seq_pos: Tensor, shape ``[batch_size, seq_len]``
    """
    pes = self.pe.expand(hidden.size(0), -1, -1).gather(
        1, seq_pos.unsqueeze(-1).expand(-1, -1, self.d_model)
    )
    hidden = hidden + pes
    return self.dropout(hidden)

RandomEncoding ¶

RandomEncoding(embed_dim: int, max_classes: int = 100)

Bases: Module

This is like torch.nn.Embedding but with rows of embeddings are randomly permuted in each forward pass before lookup operation. This might be useful in cases where classes have no fixed meaning but rather indicate a connection between different elements in a sequence. Reference is the MatNet model.

Source code in rl4co/models/nn/ops.py

def __init__(self, embed_dim: int, max_classes: int = 100):
    super().__init__()
    self.embed_dim = embed_dim
    self.max_classes = max_classes
    rand_emb = torch.rand(max_classes, self.embed_dim)
    self.register_buffer("emb", rand_emb)

Positional Encodings¶

Routing positional encodings. IPE/XPE/Hierarchical: https://arxiv.org/abs/2605.11910

Classes:

AbsolutePE –

Absolute PE (APE): a learnable lookup table indexed by within-route position.
SinusoidalPE –

Sinusoidal PE (SIN): fixed sin/cos of the integer index (Vaswani et al., 2017).
RotaryPE –

Rotary PE (RoPE): query/key rotation rather than an additive embedding.
RelativePE –

Relative PE (RPE): a learnable bias indexed by signed index offset (Shaw et al., 2018).
ALiBiBias –

ALiBi: a fixed linear penalty subtracted from attention logits (Press et al., 2021).
DACTCyclicPE –

DACT cyclic PE: a Gray-code lookup over the cyclic index, projected to D.
CycleFormerPE –

CycleFormer circular PE: a sinusoidal map of the index wrapped around the tour length.
LaplacianPE –

Laplacian PE (Lap. PE): the K smallest non-trivial Laplacian eigenvectors.
RandomWalkSE –

Random-Walk Structural Encoding (RWSE): diagonals of random-walk matrix powers.
ShortestPathBias –

Shortest-path-distance attention bias (SPD): a learnable scalar per integer distance.
InRoutePE –

In-route PE (IPE): distance-indexed multi-frequency sinusoid (paper §4.1).
CrossRoutePE –

Cross-route PE (XPE): depot-anchored polar-angle sinusoid (paper §4.2).
HierarchicalPE –

Hierarchical PE: per-node concatenation [IPE(v) ‖ XPE(v)] (paper §4.3, pre-projection).

Functions:

build_route_graph –

Build a symmetric adjacency matrix from padded route permutations.
get_positional_encoding –

Instantiate a positional-encoding module by name.

AbsolutePE ¶

AbsolutePE(embed_dim: int, max_len: int = 1000)

Bases: Module

Absolute PE (APE): a learnable lookup table indexed by within-route position.

Implements PE^APE(v_i) = E_i with E ∈ R^{(max_len + 1) × D} trainable (index 0 reserved for the depot).

Parameters:

embed_dim (int) –

Embedding dimension D.
max_len (int, default: 1000 ) –

Maximum within-route position the table can index (L_max).

Methods:

forward –

Look up the per-position embeddings.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, max_len: int = 1000) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.max_len = max_len
    self.embedding = nn.Embedding(max_len + 1, embed_dim)

forward ¶

forward(positions: Tensor) -> Tensor

Look up the per-position embeddings.

Parameters:

positions (Tensor) –

Long tensor of shape [..., L] with within-route indices.

Returns:

Tensor –

Tensor of shape [..., L, D].

Source code in rl4co/utils/pe.py

def forward(self, positions: Tensor) -> Tensor:
    """Look up the per-position embeddings.

    Args:
        positions: Long tensor of shape ``[..., L]`` with within-route indices.

    Returns:
        Tensor of shape ``[..., L, D]``.
    """
    return self.embedding(positions.long().clamp(max=self.max_len))

SinusoidalPE ¶

SinusoidalPE(embed_dim: int, base: float = 10000.0)

Bases: Module

Sinusoidal PE (SIN): fixed sin/cos of the integer index (Vaswani et al., 2017).

PE^SIN_{2k}(v_i) = sin(i / λ^{2k/D}) and PE^SIN_{2k+1}(v_i) = cos(i / λ^{2k/D}) with λ = 10000 and k = 0 .. D/2 - 1.

Parameters:

embed_dim (int) –

Embedding dimension D (assumed even; an odd D drops the last channel of the cosine half).
base (float, default: 10000.0 ) –

Wavelength base λ (default 10000).

Methods:

forward –

Compute the sinusoidal encoding for the given positions.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, base: float = 10000.0) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.register_buffer("freqs", _geometric_frequencies((embed_dim + 1) // 2, embed_dim, base))

forward ¶

forward(positions: Tensor) -> Tensor

Compute the sinusoidal encoding for the given positions.

Parameters:

positions (Tensor) –

Tensor of shape [..., L] (int or float) with within-route indices.

Returns:

Tensor –

Tensor of shape [..., L, D].

Source code in rl4co/utils/pe.py

def forward(self, positions: Tensor) -> Tensor:
    """Compute the sinusoidal encoding for the given positions.

    Args:
        positions: Tensor of shape ``[..., L]`` (int or float) with within-route indices.

    Returns:
        Tensor of shape ``[..., L, D]``.
    """
    angles = positions.unsqueeze(-1).float() * self.freqs  # [..., L, D/2]
    out = torch.zeros(
        *angles.shape[:-1], self.embed_dim, device=angles.device, dtype=angles.dtype
    )
    out[..., 0::2] = torch.sin(angles)
    out[..., 1::2] = torch.cos(angles[..., : out[..., 1::2].shape[-1]])
    return out

RotaryPE ¶

RotaryPE(embed_dim: int, base: float = 10000.0)

Bases: Module

Rotary PE (RoPE): query/key rotation rather than an additive embedding.

For each channel pair (2k, 2k+1) the query/key is rotated by θ_i = i · λ^{-2k/D} (Su et al., 2024). The resulting attention logit q_i^T k_j depends only on the index difference i - j.

Parameters:

embed_dim (int) –

Per-head dimension D (assumed even).
base (float, default: 10000.0 ) –

Wavelength base λ (default 10000).

Methods:

rotate –

Apply the rotary transform to a single tensor.
rotate_queries_keys –

Rotate both queries and keys by their positions.
forward –

Alias for :meth:rotate_queries_keys.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, base: float = 10000.0) -> None:
    super().__init__()
    assert embed_dim % 2 == 0, "RotaryPE requires an even embed_dim"
    self.embed_dim = embed_dim
    self.register_buffer("freqs", _geometric_frequencies(embed_dim // 2, embed_dim, base))

rotate ¶

rotate(x: Tensor, positions: Tensor) -> Tensor

Apply the rotary transform to a single tensor.

Parameters:

x (Tensor) –

Tensor of shape [..., L, D].
positions (Tensor) –

Tensor of shape [..., L] with the index of each element.

Returns:

Tensor –

Rotated tensor of the same shape as x.

Source code in rl4co/utils/pe.py

def rotate(self, x: Tensor, positions: Tensor) -> Tensor:
    """Apply the rotary transform to a single tensor.

    Args:
        x: Tensor of shape ``[..., L, D]``.
        positions: Tensor of shape ``[..., L]`` with the index of each element.

    Returns:
        Rotated tensor of the same shape as ``x``.
    """
    cos, sin = self._rotary_components(positions)
    return x * cos + self._rotate_half(x) * sin

rotate_queries_keys ¶

rotate_queries_keys(
    q: Tensor, k: Tensor, positions: Tensor
) -> tuple[Tensor, Tensor]

Rotate both queries and keys by their positions.

Parameters:

q (Tensor) –

Query tensor [..., L, D].
k (Tensor) –

Key tensor [..., L, D].
positions (Tensor) –

Index tensor [..., L] shared by q and k.

Returns:

tuple[Tensor, Tensor] –

Tuple (q_rot, k_rot) with the same shapes as the inputs.

Source code in rl4co/utils/pe.py

def rotate_queries_keys(self, q: Tensor, k: Tensor, positions: Tensor) -> tuple[Tensor, Tensor]:
    """Rotate both queries and keys by their positions.

    Args:
        q: Query tensor ``[..., L, D]``.
        k: Key tensor ``[..., L, D]``.
        positions: Index tensor ``[..., L]`` shared by ``q`` and ``k``.

    Returns:
        Tuple ``(q_rot, k_rot)`` with the same shapes as the inputs.
    """
    return self.rotate(q, positions), self.rotate(k, positions)

forward ¶

forward(
    q: Tensor, k: Tensor, positions: Tensor
) -> tuple[Tensor, Tensor]

Alias for :meth:rotate_queries_keys.

Source code in rl4co/utils/pe.py

def forward(self, q: Tensor, k: Tensor, positions: Tensor) -> tuple[Tensor, Tensor]:
    """Alias for :meth:`rotate_queries_keys`."""
    return self.rotate_queries_keys(q, k, positions)

RelativePE ¶

RelativePE(window: int = 16)

Bases: Module

Relative PE (RPE): a learnable bias indexed by signed index offset (Shaw et al., 2018).

logit(i, j) += b_{clip(i - j, -W, W)} with {b_Δ}_{Δ=-W}^{W} shared across heads.

Parameters:

window (int, default: 16 ) –

Clipping window W; offsets are clipped to [-W, W].

Methods:

forward –

Build the [L, L] bias matrix added to attention logits.

Source code in rl4co/utils/pe.py

def __init__(self, window: int = 16) -> None:
    super().__init__()
    self.window = window
    self.bias = nn.Embedding(2 * window + 1, 1)

forward ¶

forward(
    seq_len: int, device: device | None = None
) -> Tensor

Build the [L, L] bias matrix added to attention logits.

Parameters:

seq_len (int) –

Sequence length L.
device (device | None, default: None ) –

Optional device for the output.

Returns:

Tensor –

Tensor of shape [L, L] where entry (i, j) is b_{clip(i-j, -W, W)}.

Source code in rl4co/utils/pe.py

def forward(self, seq_len: int, device: torch.device | None = None) -> Tensor:
    """Build the ``[L, L]`` bias matrix added to attention logits.

    Args:
        seq_len: Sequence length ``L``.
        device: Optional device for the output.

    Returns:
        Tensor of shape ``[L, L]`` where entry ``(i, j)`` is ``b_{clip(i-j, -W, W)}``.
    """
    idx = torch.arange(seq_len, device=device)
    offset = (idx[:, None] - idx[None, :]).clamp(-self.window, self.window) + self.window
    return self.bias(offset).squeeze(-1)

ALiBiBias ¶

ALiBiBias(num_heads: int)

Bases: Module

ALiBi: a fixed linear penalty subtracted from attention logits (Press et al., 2021).

logit(i, j) -= m_h · |i - j| with head-specific slopes m_h following the geometric schedule of the original paper. The returned tensor is shaped [num_heads, L, L] (broadcasts over the batch dimension).

Parameters:

num_heads (int) –

Number of attention heads n_h.

Methods:

forward –

Build the [num_heads, L, L] ALiBi bias.

Source code in rl4co/utils/pe.py

def __init__(self, num_heads: int) -> None:
    super().__init__()
    self.num_heads = num_heads
    self.register_buffer("slopes", self._get_slopes(num_heads))

forward ¶

forward(
    seq_len: int, device: device | None = None
) -> Tensor

Build the [num_heads, L, L] ALiBi bias.

Parameters:

seq_len (int) –

Sequence length L.
device (device | None, default: None ) –

Optional device for the output.

Returns:

Tensor –

Tensor of shape [num_heads, L, L] equal to -m_h · |i - j|.

Source code in rl4co/utils/pe.py

def forward(self, seq_len: int, device: torch.device | None = None) -> Tensor:
    """Build the ``[num_heads, L, L]`` ALiBi bias.

    Args:
        seq_len: Sequence length ``L``.
        device: Optional device for the output.

    Returns:
        Tensor of shape ``[num_heads, L, L]`` equal to ``-m_h · |i - j|``.
    """
    idx = torch.arange(seq_len, device=device)
    dist = (idx[:, None] - idx[None, :]).abs().float()  # [L, L]
    slopes = self.slopes.to(device=device)
    return -slopes[:, None, None] * dist[None, :, :]

DACTCyclicPE ¶

DACTCyclicPE(embed_dim: int, max_len: int = 1000)

Bases: Module

DACT cyclic PE: a Gray-code lookup over the cyclic index, projected to D.

PE^DACT(v_i) = Linear(Gray(i mod L)) (Ma et al., 2021). The Gray code uses ceil(log2(L)) bits so consecutive cyclic positions differ by exactly one bit; this is not bit-wise circular when L is not a power of two (documented limitation).

Parameters:

embed_dim (int) –

Output embedding dimension D.
max_len (int, default: 1000 ) –

Maximum route length L supported (sets the number of Gray-code bits).

Methods:

forward –

Encode the within-route positions.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, max_len: int = 1000) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.max_len = max_len
    self.num_bits = max(1, math.ceil(math.log2(max(2, max_len))))
    codes = torch.tensor(
        [_binary_reflected_gray(n, self.num_bits) for n in range(max_len)], dtype=torch.float32
    )
    self.register_buffer("gray_codes", codes)  # [max_len, num_bits]
    self.proj = nn.Linear(self.num_bits, embed_dim)

forward ¶

forward(
    positions: Tensor, seq_len: int | None = None
) -> Tensor

Encode the within-route positions.

Parameters:

positions (Tensor) –

Long tensor [..., L] of within-route indices.
seq_len (int | None, default: None ) –

Route length L used for the cyclic i mod L reduction. Defaults to positions.shape[-1].

Returns:

Tensor –

Tensor of shape [..., L, D].

Source code in rl4co/utils/pe.py

def forward(self, positions: Tensor, seq_len: int | None = None) -> Tensor:
    """Encode the within-route positions.

    Args:
        positions: Long tensor ``[..., L]`` of within-route indices.
        seq_len: Route length ``L`` used for the cyclic ``i mod L`` reduction.  Defaults
            to ``positions.shape[-1]``.

    Returns:
        Tensor of shape ``[..., L, D]``.
    """
    if seq_len is None:
        seq_len = positions.shape[-1]
    idx = positions.long() % max(1, seq_len)
    codes = self.gray_codes[idx.clamp(max=self.max_len - 1)]  # [..., L, num_bits]
    return self.proj(codes)

CycleFormerPE ¶

CycleFormerPE(embed_dim: int, base: float = 10000.0)

Bases: Module

CycleFormer circular PE: a sinusoidal map of the index wrapped around the tour length.

PE^Cyc_{2k}(v_i) = sin(2π (i mod L) / L · ω_k) and the cosine counterpart, with the same geometric frequency schedule as :class:SinusoidalPE (Yook et al., 2024). Reducing the index modulo L makes the encoding identical for i and i + L regardless of the frequency schedule.

Parameters:

embed_dim (int) –

Embedding dimension D (assumed even).
base (float, default: 10000.0 ) –

Wavelength base λ for the geometric schedule (default 10000).

Methods:

forward –

Compute the circular sinusoidal encoding.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, base: float = 10000.0) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.register_buffer("freqs", _geometric_frequencies((embed_dim + 1) // 2, embed_dim, base))

forward ¶

forward(
    positions: Tensor, seq_len: int | None = None
) -> Tensor

Compute the circular sinusoidal encoding.

Parameters:

positions (Tensor) –

Tensor [..., L] of within-route indices.
seq_len (int | None, default: None ) –

Tour length L for the 2π i / L wrap. Defaults to positions.shape[-1].

Returns:

Tensor –

Tensor of shape [..., L, D].

Source code in rl4co/utils/pe.py

def forward(self, positions: Tensor, seq_len: int | None = None) -> Tensor:
    """Compute the circular sinusoidal encoding.

    Args:
        positions: Tensor ``[..., L]`` of within-route indices.
        seq_len: Tour length ``L`` for the ``2π i / L`` wrap.  Defaults to
            ``positions.shape[-1]``.

    Returns:
        Tensor of shape ``[..., L, D]``.
    """
    if seq_len is None:
        seq_len = positions.shape[-1]
    idx = positions.float() % float(max(1, seq_len))
    phase = 2.0 * math.pi * idx / float(max(1, seq_len))  # [..., L]
    angles = phase.unsqueeze(-1) * self.freqs  # [..., L, D/2]
    out = torch.zeros(
        *angles.shape[:-1], self.embed_dim, device=angles.device, dtype=angles.dtype
    )
    out[..., 0::2] = torch.sin(angles)
    out[..., 1::2] = torch.cos(angles[..., : out[..., 1::2].shape[-1]])
    return out

LaplacianPE ¶

LaplacianPE(embed_dim: int, k: int = 8)

Bases: Module

Laplacian PE (Lap. PE): the K smallest non-trivial Laplacian eigenvectors.

For each instance the route graph is built (depot plus route arcs), the unnormalized Laplacian L = diag(deg) - A is formed, and the eigenvectors of the K smallest non-trivial eigenvalues are taken (Dwivedi & Bresson, 2020). Sign ambiguity is resolved by random flips during training (deterministic, no flips, in eval). The result is zero-padded to D columns.

Note: only K_eff = min(K, N - 1) columns can be genuinely non-trivial for a graph on N nodes; the remaining D - K_eff columns are zeros. Only eigenvector 0 is dropped (as in PyG), so on a disconnected graph (e.g. unvisited nodes) the leading columns can be component indicators.

Parameters:

embed_dim (int) –

Output embedding dimension D.
k (int, default: 8 ) –

Number of non-trivial eigenvectors to use (default 8).

Methods:

forward –

Compute the Laplacian eigenvector PE.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, k: int = 8) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.k = k

forward ¶

forward(
    routes: Tensor | None = None,
    adj: Tensor | None = None,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor

Compute the Laplacian eigenvector PE.

Parameters:

routes (Tensor | None, default: None ) –

Padded route permutations [..., R, Lr] or [..., Lr].
adj (Tensor | None, default: None ) –

Alternatively, a pre-built symmetric adjacency [..., N, N].
num_nodes (int | None, default: None ) –

Number of nodes (passed to :func:build_route_graph).
pad_value (int | None, default: None ) –

Padding value to ignore (passed to :func:build_route_graph).

Returns:

Tensor –

Tensor of shape [..., N, D].

Source code in rl4co/utils/pe.py

def forward(
    self,
    routes: Tensor | None = None,
    adj: Tensor | None = None,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor:
    """Compute the Laplacian eigenvector PE.

    Args:
        routes: Padded route permutations ``[..., R, Lr]`` or ``[..., Lr]``.
        adj: Alternatively, a pre-built symmetric adjacency ``[..., N, N]``.
        num_nodes: Number of nodes (passed to :func:`build_route_graph`).
        pad_value: Padding value to ignore (passed to :func:`build_route_graph`).

    Returns:
        Tensor of shape ``[..., N, D]``.
    """
    adj = _normalize_routes_or_adj(routes, adj, num_nodes, pad_value)
    n = adj.shape[-1]
    deg = adj.sum(-1)  # [..., N]
    lap = torch.diag_embed(deg) - adj  # [..., N, N]
    # eigh returns ascending eigenvalues
    _, eigvecs = torch.linalg.eigh(lap)  # eigvecs: [..., N, N]
    k_eff = min(self.k, max(0, n - 1))
    out = torch.zeros(*adj.shape[:-1], self.embed_dim, device=adj.device, dtype=eigvecs.dtype)
    if k_eff == 0:
        return out
    # drop the first (trivial) eigenvector, take the next k_eff
    vecs = eigvecs[..., 1 : 1 + k_eff]  # [..., N, k_eff]
    if self.training:
        signs = (
            torch.randint(0, 2, vecs.shape[:-2] + (1, k_eff), device=vecs.device).float() * 2
            - 1
        )
        vecs = vecs * signs
    out[..., :k_eff] = vecs
    return out

RandomWalkSE ¶

RandomWalkSE(embed_dim: int, k: int = 8)

Bases: Module

Random-Walk Structural Encoding (RWSE): diagonals of random-walk matrix powers.

PE^RWSE(v) = [(R^k)_{vv}]_{k=1}^{K} with R = D^{-1} A on the route graph (Dwivedi et al., 2021), zero-padded to D. Isolated nodes (degree 0) get a zero row of R so their landing probabilities are 0.

Parameters:

embed_dim (int) –

Output embedding dimension D.
k (int, default: 8 ) –

Number of random-walk steps K (default 8).

Methods:

forward –

Compute the random-walk structural encoding.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, k: int = 8) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.k = k

forward ¶

forward(
    routes: Tensor | None = None,
    adj: Tensor | None = None,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor

Compute the random-walk structural encoding.

Parameters:

routes (Tensor | None, default: None ) –

Padded route permutations [..., R, Lr] or [..., Lr].
adj (Tensor | None, default: None ) –

Alternatively, a pre-built symmetric adjacency [..., N, N].
num_nodes (int | None, default: None ) –

Number of nodes (passed to :func:build_route_graph).
pad_value (int | None, default: None ) –

Padding value to ignore (passed to :func:build_route_graph).

Returns:

Tensor –

Tensor of shape [..., N, D] with entries in [0, 1].

Source code in rl4co/utils/pe.py

def forward(
    self,
    routes: Tensor | None = None,
    adj: Tensor | None = None,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor:
    """Compute the random-walk structural encoding.

    Args:
        routes: Padded route permutations ``[..., R, Lr]`` or ``[..., Lr]``.
        adj: Alternatively, a pre-built symmetric adjacency ``[..., N, N]``.
        num_nodes: Number of nodes (passed to :func:`build_route_graph`).
        pad_value: Padding value to ignore (passed to :func:`build_route_graph`).

    Returns:
        Tensor of shape ``[..., N, D]`` with entries in ``[0, 1]``.
    """
    adj = _normalize_routes_or_adj(routes, adj, num_nodes, pad_value)
    n = adj.shape[-1]
    deg = adj.sum(-1, keepdim=True).clamp_min(1e-12)  # [..., N, 1]
    rw = adj / deg  # row-stochastic where degree > 0
    out = torch.zeros(*adj.shape[:-1], self.embed_dim, device=adj.device, dtype=adj.dtype)
    kmax = min(self.k, self.embed_dim)
    if kmax == 0 or n == 0:
        return out
    power = torch.eye(n, device=adj.device, dtype=adj.dtype).expand_as(adj).clone()
    for j in range(kmax):
        power = power @ rw
        out[..., j] = torch.diagonal(power, dim1=-2, dim2=-1)
    return out.clamp(0.0, 1.0)

ShortestPathBias ¶

ShortestPathBias(max_spd: int = 16)

Bases: Module

Shortest-path-distance attention bias (SPD): a learnable scalar per integer distance.

logit(i, j) += b_{spd(v_i, v_j)} with a learnable scalar bias per integer shortest-path distance on the route graph (Ying et al., 2021), capped at max_spd (unreachable pairs are clamped to max_spd).

Parameters:

max_spd (int, default: 16 ) –

Maximum shortest-path distance the bias table indexes (default 16).

Methods:

forward –

Compute the [..., N, N] SPD bias added to attention logits.

Source code in rl4co/utils/pe.py

def __init__(self, max_spd: int = 16) -> None:
    super().__init__()
    self.max_spd = max_spd
    self.bias = nn.Embedding(max_spd + 1, 1)

forward ¶

forward(
    routes: Tensor | None = None,
    adj: Tensor | None = None,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor

Compute the [..., N, N] SPD bias added to attention logits.

Parameters:

routes (Tensor | None, default: None ) –

Padded route permutations [..., R, Lr] or [..., Lr].
adj (Tensor | None, default: None ) –

Alternatively, a pre-built symmetric adjacency [..., N, N].
num_nodes (int | None, default: None ) –

Number of nodes (passed to :func:build_route_graph).
pad_value (int | None, default: None ) –

Padding value to ignore (passed to :func:build_route_graph).

Returns:

Tensor –

Tensor of shape [..., N, N]. It is symmetric and the spd = 0 (diagonal)
Tensor –

entries all share the single learnable bias b_0.

Source code in rl4co/utils/pe.py

def forward(
    self,
    routes: Tensor | None = None,
    adj: Tensor | None = None,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor:
    """Compute the ``[..., N, N]`` SPD bias added to attention logits.

    Args:
        routes: Padded route permutations ``[..., R, Lr]`` or ``[..., Lr]``.
        adj: Alternatively, a pre-built symmetric adjacency ``[..., N, N]``.
        num_nodes: Number of nodes (passed to :func:`build_route_graph`).
        pad_value: Padding value to ignore (passed to :func:`build_route_graph`).

    Returns:
        Tensor of shape ``[..., N, N]``.  It is symmetric and the ``spd = 0`` (diagonal)
        entries all share the single learnable bias ``b_0``.
    """
    adj = _normalize_routes_or_adj(routes, adj, num_nodes, pad_value)
    dist = self._all_pairs_shortest_paths(adj)
    idx = dist.clamp(0, self.max_spd).round().long()
    return self.bias(idx).squeeze(-1)

InRoutePE ¶

InRoutePE(
    embed_dim: int,
    direction_aware: bool = False,
    eps: float = 1e-09,
)

Bases: Module

In-route PE (IPE): distance-indexed multi-frequency sinusoid (paper §4.1).

For a closed route r = (v_1, ..., v_L) with v_1 = v_L = depot, define the cumulative travel distance d_i = Σ_{j=2}^{i} ||x_{v_j} - x_{v_{j-1}}|| and rescale it to one period d̂_i = 2π d_i / d_L ∈ [0, 2π]. Then

direction_aware=True → [sin(ω_k d̂_i), cos(ω_k d̂_i)]_{k=0}^{D/2-1} (distinguishes a route from its reversal: reversal maps d_i → d_L - d_i, flipping the sine half).
direction_aware=False → [cos(ω_k d̂_i)]_{k=0}^{D-1} (invariant to route reversal by construction, since cos(2π - x) = cos(x) with integer frequencies).

Implementation note: the paper writes ω_k = λ^{-2k/D} (geometric), but the stated properties IPE(v_1) = IPE(v_L) and the reversal (in)variance only hold for integer frequencies. This module therefore uses integer harmonics ω_k = k + 1.

Parameters:

embed_dim (int) –

Output embedding dimension D.
direction_aware (bool, default: False ) –

Whether to use the sin/cos variant (True) or the cosine-only reversal-invariant variant (False).
eps (float, default: 1e-09 ) –

Small constant guarding the d_L division for degenerate (zero-length) routes.

Methods:

forward –

Compute the in-route PE for route-ordered coordinates.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, direction_aware: bool = False, eps: float = 1e-9) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.direction_aware = direction_aware
    self.eps = eps
    if direction_aware:
        num_freq = (embed_dim + 1) // 2
    else:
        num_freq = embed_dim
    self.register_buffer("freqs", torch.arange(1, num_freq + 1, dtype=torch.float32))

forward ¶

forward(coords: Tensor) -> Tensor

Compute the in-route PE for route-ordered coordinates.

Parameters:

coords (Tensor) –

Route-ordered coordinates of shape [..., L, 2]. The route is assumed closed (coords[..., 0, :] == coords[..., -1, :] == depot).

Returns:

Tensor –

Tensor of shape [..., L, D].

Source code in rl4co/utils/pe.py

def forward(self, coords: Tensor) -> Tensor:
    """Compute the in-route PE for route-ordered coordinates.

    Args:
        coords: Route-ordered coordinates of shape ``[..., L, 2]``.  The route is assumed
            closed (``coords[..., 0, :] == coords[..., -1, :] == depot``).

    Returns:
        Tensor of shape ``[..., L, D]``.
    """
    seg = torch.linalg.norm(coords[..., 1:, :] - coords[..., :-1, :], dim=-1)  # [..., L-1]
    cum = torch.cumsum(seg, dim=-1)  # [..., L-1]
    zero = torch.zeros(*cum.shape[:-1], 1, device=cum.device, dtype=cum.dtype)
    d = torch.cat([zero, cum], dim=-1)  # [..., L], d_1 = 0
    d_total = d[..., -1:].clamp_min(self.eps)  # [..., 1]
    d_hat = 2.0 * math.pi * d / d_total  # [..., L]
    angles = d_hat.unsqueeze(-1) * self.freqs  # [..., L, num_freq]
    if self.direction_aware:
        out = torch.zeros(
            *angles.shape[:-1], self.embed_dim, device=angles.device, dtype=angles.dtype
        )
        out[..., 0::2] = torch.sin(angles)
        out[..., 1::2] = torch.cos(angles[..., : out[..., 1::2].shape[-1]])
        return out
    return torch.cos(angles)  # [..., L, D]

CrossRoutePE ¶

CrossRoutePE(embed_dim: int, k: int = 4)

Bases: Module

Cross-route PE (XPE): depot-anchored polar-angle sinusoid (paper §4.2).

For node v with coordinate x_v and depot x_0, define θ_v = atan2(y_v - y_0, x_v - x_0) ∈ [-π, π) and encode

XPE(v) = [sin(ω'_k θ_v), cos(ω'_k θ_v)]_{k=0}^{K-1}, ω'_k = 2^k,

zero-padded from 2K to D columns (requires 2K ≤ D). A customer coincident with the depot yields atan2(0, 0) = 0 (acceptable).

Parameters:

embed_dim (int) –

Output embedding dimension D.
k (int, default: 4 ) –

Number of frequency bands K (default 4; must satisfy 2K ≤ D).

Methods:

forward –

Compute the cross-route PE.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, k: int = 4) -> None:
    super().__init__()
    if 2 * k > embed_dim:
        raise ValueError(
            f"CrossRoutePE requires 2*k <= embed_dim, got k={k}, embed_dim={embed_dim}"
        )
    self.embed_dim = embed_dim
    self.k = k
    self.register_buffer("freqs", 2.0 ** torch.arange(k, dtype=torch.float32))

forward ¶

forward(coords: Tensor, depot: Tensor) -> Tensor

Compute the cross-route PE.

Parameters:

coords (Tensor) –

Node coordinates of shape [..., 2] or [..., L, 2].
depot (Tensor) –

Depot coordinate of shape [..., 2] (broadcast against coords).

Returns:

Tensor –

Tensor of shape coords.shape[:-1] + (D,).

Source code in rl4co/utils/pe.py

def forward(self, coords: Tensor, depot: Tensor) -> Tensor:
    """Compute the cross-route PE.

    Args:
        coords: Node coordinates of shape ``[..., 2]`` or ``[..., L, 2]``.
        depot: Depot coordinate of shape ``[..., 2]`` (broadcast against ``coords``).

    Returns:
        Tensor of shape ``coords.shape[:-1] + (D,)``.
    """
    depot_b = depot
    while depot_b.dim() < coords.dim():
        depot_b = depot_b.unsqueeze(-2)
    rel = coords - depot_b  # [..., 2]
    theta = torch.atan2(rel[..., 1], rel[..., 0])  # [...]
    angles = theta.unsqueeze(-1) * self.freqs  # [..., K]
    out = torch.zeros(*theta.shape, self.embed_dim, device=coords.device, dtype=coords.dtype)
    out[..., 0 : 2 * self.k : 2] = torch.sin(angles)
    out[..., 1 : 2 * self.k : 2] = torch.cos(angles)
    return out

HierarchicalPE ¶

HierarchicalPE(
    embed_dim: int,
    direction_aware: bool = False,
    xpe_k: int = 4,
)

Bases: Module

Hierarchical PE: per-node concatenation [IPE(v) ‖ XPE(v)] (paper §4.3, pre-projection).

This produces the positional part of the §4.3 fusion FF([x_v ‖ IPE(v) ‖ XPE(v)]); the raw-coordinate concatenation and the shared feed-forward projection itself are out of scope for this module. The output width is 2 * embed_dim (IPE width embed_dim plus XPE width embed_dim).

Parameters:

embed_dim (int) –

Per-component embedding dimension D; the output width is 2D.
direction_aware (bool, default: False ) –

Forwarded to :class:InRoutePE.
xpe_k (int, default: 4 ) –

Forwarded to :class:CrossRoutePE as k (default 4).

Methods:

forward –

Compute the concatenated hierarchical PE for a route.

Source code in rl4co/utils/pe.py

def __init__(self, embed_dim: int, direction_aware: bool = False, xpe_k: int = 4) -> None:
    super().__init__()
    self.embed_dim = embed_dim
    self.out_dim = 2 * embed_dim
    self.ipe = InRoutePE(embed_dim, direction_aware=direction_aware)
    self.xpe = CrossRoutePE(embed_dim, k=xpe_k)

forward ¶

forward(
    coords: Tensor, depot: Tensor | None = None
) -> Tensor

Compute the concatenated hierarchical PE for a route.

Parameters:

coords (Tensor) –

Route-ordered coordinates [..., L, 2] (closed route; index 0 is the depot).
depot (Tensor | None, default: None ) –

Optional depot coordinate [..., 2]. Defaults to coords[..., 0, :].

Returns:

Tensor –

Tensor of shape [..., L, 2D].

Source code in rl4co/utils/pe.py

def forward(self, coords: Tensor, depot: Tensor | None = None) -> Tensor:
    """Compute the concatenated hierarchical PE for a route.

    Args:
        coords: Route-ordered coordinates ``[..., L, 2]`` (closed route; index ``0`` is
            the depot).
        depot: Optional depot coordinate ``[..., 2]``.  Defaults to ``coords[..., 0, :]``.

    Returns:
        Tensor of shape ``[..., L, 2D]``.
    """
    if depot is None:
        depot = coords[..., 0, :]
    ipe = self.ipe(coords)  # [..., L, D]
    xpe = self.xpe(coords, depot)  # [..., L, D]
    return torch.cat([ipe, xpe], dim=-1)

build_route_graph ¶

build_route_graph(
    routes: Tensor,
    num_nodes: int | None = None,
    pad_value: int | None = None,
) -> Tensor

Build a symmetric adjacency matrix from padded route permutations.

Edges connect consecutive nodes on each route; the depot (index 0) connects to the route endpoints through the usual closed representation [0, c_1, ..., c_k, 0, ...]. Self-loops (e.g. from zero-padding [..., c_k, 0, 0, 0]) are dropped, and any arc touching pad_value (when given) is ignored.

Parameters:

routes (Tensor) –

Long tensor of shape [..., R, Lr] (multiple routes per instance) or [..., Lr] (a single route per instance).
num_nodes (int | None, default: None ) –

Number of nodes N (including the depot). Defaults to routes.max() + 1.
pad_value (int | None, default: None ) –

Optional padding value to ignore (e.g. -1). If None, only self-loops are removed.

Returns:

Tensor –

Symmetric adjacency tensor of shape [..., N, N] with {0, 1} entries and a
Tensor –

zero diagonal.

Source code in rl4co/utils/pe.py

def build_route_graph(
    routes: Tensor, num_nodes: int | None = None, pad_value: int | None = None
) -> Tensor:
    """Build a symmetric adjacency matrix from padded route permutations.

    Edges connect consecutive nodes on each route; the depot (index ``0``) connects to the
    route endpoints through the usual closed representation ``[0, c_1, ..., c_k, 0, ...]``.
    Self-loops (e.g. from zero-padding ``[..., c_k, 0, 0, 0]``) are dropped, and any arc
    touching ``pad_value`` (when given) is ignored.

    Args:
        routes: Long tensor of shape ``[..., R, Lr]`` (multiple routes per instance) or
            ``[..., Lr]`` (a single route per instance).
        num_nodes: Number of nodes ``N`` (including the depot).  Defaults to
            ``routes.max() + 1``.
        pad_value: Optional padding value to ignore (e.g. ``-1``).  If ``None``, only
            self-loops are removed.

    Returns:
        Symmetric adjacency tensor of shape ``[..., N, N]`` with ``{0, 1}`` entries and a
        zero diagonal.
    """
    routes = routes.long()
    # Normalize to shape [..., R, Lr] (a route axis is inserted for [..., Lr] inputs).
    if routes.dim() == 1:
        routes = routes.view(1, 1, -1)
    elif routes.dim() == 2:
        # Ambiguous [R, Lr] vs [B, Lr]; treat as [B, Lr] -> [B, 1, Lr]. A bare [R, Lr] still
        # yields the right graph since edges from all R rows are unioned anyway.
        routes = routes.unsqueeze(-2)
    # else: already [..., R, Lr]

    if num_nodes is None:
        num_nodes = int(routes.max().item()) + 1 if routes.numel() > 0 else 1
        if pad_value is not None and num_nodes - 1 == pad_value:
            num_nodes = int(routes[routes != pad_value].max().item()) + 1

    *batch, n_routes, route_len = routes.shape
    adj = torch.zeros(*batch, num_nodes, num_nodes, device=routes.device)
    if route_len < 2:
        return adj

    src = routes[..., :-1].reshape(*batch, -1)  # [..., R*(Lr-1)]
    dst = routes[..., 1:].reshape(*batch, -1)
    valid = src != dst
    if pad_value is not None:
        valid = valid & (src != pad_value) & (dst != pad_value)
    valid = valid & (src >= 0) & (dst >= 0) & (src < num_nodes) & (dst < num_nodes)

    src = src.clamp(0, num_nodes - 1)
    dst = dst.clamp(0, num_nodes - 1)
    w = valid.float()
    # Scatter into the adjacency for each batch element.
    flat_adj = adj.reshape(-1, num_nodes, num_nodes)
    flat_src = src.reshape(-1, src.shape[-1])
    flat_dst = dst.reshape(-1, dst.shape[-1])
    flat_w = w.reshape(-1, w.shape[-1])
    for b in range(flat_adj.shape[0]):
        flat_adj[b].index_put_((flat_src[b], flat_dst[b]), flat_w[b], accumulate=True)
        flat_adj[b].index_put_((flat_dst[b], flat_src[b]), flat_w[b], accumulate=True)
    adj = flat_adj.reshape(*batch, num_nodes, num_nodes)
    adj = (adj > 0).float()
    # zero the diagonal
    eye = torch.eye(num_nodes, device=adj.device, dtype=adj.dtype)
    adj = adj * (1.0 - eye)
    return adj

get_positional_encoding ¶

get_positional_encoding(name: str, **kwargs) -> Module

Instantiate a positional-encoding module by name.

Parameters:

name (str) –

One of "APE" | "SIN" | "RoPE" | "RPE" | "ALiBi" | "LapPE" | "RWSE" | "SPD" | "DACT" | "CycleFormer" | "IPE" | "XPE" | "Hierarchical".
**kwargs –

Forwarded to the selected class constructor.

Returns:

Module –

The instantiated :class:torch.nn.Module.

Raises:

ValueError –

If name is not a registered positional encoding.

Source code in rl4co/utils/pe.py

def get_positional_encoding(name: str, **kwargs) -> nn.Module:
    """Instantiate a positional-encoding module by name.

    Args:
        name: One of ``"APE" | "SIN" | "RoPE" | "RPE" | "ALiBi" | "LapPE" | "RWSE" | "SPD" |
            "DACT" | "CycleFormer" | "IPE" | "XPE" | "Hierarchical"``.
        **kwargs: Forwarded to the selected class constructor.

    Returns:
        The instantiated :class:`torch.nn.Module`.

    Raises:
        ValueError: If ``name`` is not a registered positional encoding.
    """
    if name not in _PE_REGISTRY:
        raise ValueError(
            f"Unknown positional encoding '{name}'. Available positional encodings: "
            f"{list(_PE_REGISTRY.keys())}"
        )
    return _PE_REGISTRY[name](**kwargs)