Search

class rl4co.models.zoo.common.search.base.SearchBase(env, policy, dataset, batch_size=1, max_iters=100, max_runtime=86400, save_path=None, **kwargs)

Bases: RL4COLitModule

Base class for search algorithms. Search algorithms are used onlin to find better solutions for a given dataset, i.e. given a policy, improve (a part of) its parameters such that the policy performs better on the given dataset.

Note

By default, we use manual optimization to handle the search.

Parameters:

• env – RL4CO environment

• policy – policy network

• dataset (Union[Dataset, str]) – dataset to use for training

• batch_size (int) – batch size

• **kwargs – additional arguments

on_train_batch_end(outputs, batch, batch_idx)

Called when the train batch ends. This can be used for instance for logging or clearing cache.

Return type:

None

on_train_batch_start(batch, batch_idx)

Called before training (i.e. search) for a new batch begins. This can be used to perform changes to the model or optimizer at the start of each batch.

on_train_epoch_end()

Called when the train ends.

Return type:

None

setup(stage='fit')

Setup the dataset and attributes. The RL4COLitModulebase class automatically loads the data.

test_step(batch, batch_idx)

Not used during search

training_step(batch, batch_idx)

Main search loop. We use the training step to effectively adapt to a batch of instances.

validation_step(batch, batch_idx)

Not used during search