Environment Embeddings

In autoregressive policies, environment embeddings transfer data from feature space to hidden space:

• Initial Embeddings: encode global problem features

• Context Embeddings: modify current node embedding during decoding

• Dynamic Embeddings: modify all nodes embeddings during decoding

Context Embeddings

The context embedding is used to modify the query embedding of the problem node of the current partial solution. Usually consists of a projection of gathered node embeddings and features to the embedding space.

class rl4co.models.nn.env_embeddings.context.DPPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Decap Placement Problem (DPP), EDA (electronic design automation). Project the following to the embedding space:

• current cell embedding

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(embeddings, td)

Context cannot be defined by a single node embedding for DPP, hence 0. We modify the dynamic embedding instead to capture placed items

class rl4co.models.nn.env_embeddings.context.EnvContext(embedding_dim, step_context_dim=None, linear_bias=False)

Bases: Module

Base class for environment context embeddings. The context embedding is used to modify the query embedding of the problem node of the current partial solution. Consists of a linear layer that projects the node features to the embedding space.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(embeddings, td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.context.MTSPContext(embedding_dim, linear_bias=False)

Bases: EnvContext

Context embedding for the Multiple Traveling Salesman Problem (mTSP). Project the following to the embedding space:

• current node embedding

• remaining_agents

• current_length

• max_subtour_length

• distance_from_depot

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class rl4co.models.nn.env_embeddings.context.OPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Orienteering Problem (OP). Project the following to the embedding space:

• current node embedding

• remaining distance (max_length - tour_length)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class rl4co.models.nn.env_embeddings.context.PCTSPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Prize Collecting TSP (PCTSP). Project the following to the embedding space:

• current node embedding

• remaining prize (prize_required - cur_total_prize)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class rl4co.models.nn.env_embeddings.context.PDPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Pickup and Delivery Problem (PDP). Project the following to the embedding space:

• current node embedding

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(embeddings, td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.context.SMTWTPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Single Machine Total Weighted Tardiness Problem (SMTWTP). Project the following to the embedding space:

• current node embedding

• current time

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class rl4co.models.nn.env_embeddings.context.TSPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Traveling Salesman Problem (TSP). Project the following to the embedding space:

• first node embedding

• current node embedding

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(embeddings, td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.context.VRPContext(embedding_dim)

Bases: EnvContext

Context embedding for the Capacitated Vehicle Routing Problem (CVRP). Project the following to the embedding space:

• current node embedding

• remaining capacity (vehicle_capacity - used_capacity)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

rl4co.models.nn.env_embeddings.context.env_context_embedding(env_name, config)

Get environment context embedding. The context embedding is used to modify the query embedding of the problem node of the current partial solution. Usually consists of a projection of gathered node embeddings and features to the embedding space.

Parameters:

• env – Environment or its name.

• config (dict) – A dictionary of configuration options for the environment.

Return type:

Module

---

Dynamic Embeddings

The dynamic embedding is used to modify query, key and value vectors of the attention mechanism based on the current state of the environment (which is changing during the rollout). Generally consists of a linear layer that projects the node features to the embedding space.

class rl4co.models.nn.env_embeddings.dynamic.SDVRPDynamicEmbedding(embedding_dim, linear_bias=False)

Bases: Module

Dynamic embedding for the Split Delivery Vehicle Routing Problem (SDVRP). Embed the following node features to the embedding space:

• demand_with_depot: demand of the customers and the depot

The demand with depot is used to modify the query, key and value vectors of the attention mechanism based on the current state of the environment (which is changing during the rollout).

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.dynamic.StaticEmbedding(*args, **kwargs)

Bases: Module

Static embedding for general problems. This is used for problems that do not have any dynamic information, except for the information regarding the current action (e.g. the current node in TSP). See context embedding for more details.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

rl4co.models.nn.env_embeddings.dynamic.env_dynamic_embedding(env_name, config)

Get environment dynamic embedding. The dynamic embedding is used to modify query, key and value vectors of the attention mechanism based on the current state of the environment (which is changing during the rollout). Consists of a linear layer that projects the node features to the embedding space.

Parameters:

• env – Environment or its name.

• config (dict) – A dictionary of configuration options for the environment.

Return type:

Module

---

Init Embeddings

The init embedding is used to initialize the general embedding of the problem nodes without any solution information. Generally consists of a linear layer that projects the node features to the embedding space.

class rl4co.models.nn.env_embeddings.init.DPPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Decap Placement Problem (DPP), EDA (electronic design automation). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (cells)

• probe: index of the (single) probe cell. We embed the euclidean distance from the probe to all cells.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.MDPPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Multi-port Placement Problem (MDPP), EDA (electronic design automation). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (cells)

• probe: indexes of the probe cells (multiple). We embed the euclidean distance of each cell to the closest probe.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.MTSPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Multiple Traveling Salesman Problem (mTSP). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (depot, cities)

NOTE: new made by Fede. May need to be checked

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.OPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Orienteering Problems (OP). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (depot and customers separately)

• prize: prize for visiting the customers

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.PCTSPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Prize Collecting Traveling Salesman Problems (PCTSP). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (depot and customers separately)

• expected_prize: expected prize for visiting the customers.

  In PCTSP, this is the actual prize. In SPCTSP, this is the expected prize.

• penalty: penalty for not visiting the customers

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.PDPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Pickup and Delivery Problem (PDP). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (depot, pickups and deliveries separately)

  Note that pickups and deliveries are interleaved in the input.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.SMTWTPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Single Machine Total Weighted Tardiness Problem (SMTWTP). Embed the following node features to the embedding space:

• job_due_time: due time of the jobs

• job_weight: weights of the jobs

• job_process_time: the processing time of jobs

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.TSPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Traveling Salesman Problems (TSP). Embed the following node features to the embedding space:

• locs: x, y coordinates of the cities

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class rl4co.models.nn.env_embeddings.init.VRPInitEmbedding(embedding_dim, linear_bias=True)

Bases: Module

Initial embedding for the Vehicle Routing Problems (VRP). Embed the following node features to the embedding space:

• locs: x, y coordinates of the nodes (depot and customers separately)

• demand: demand of the customers

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(td)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

rl4co.models.nn.env_embeddings.init.env_init_embedding(env_name, config)

Get environment initial embedding. The init embedding is used to initialize the general embedding of the problem nodes without any solution information. Consists of a linear layer that projects the node features to the embedding space.

Parameters:

• env – Environment or its name.

• config (dict) – A dictionary of configuration options for the environment.

Return type:

Module