Data

Datasets

class rl4co.data.dataset.ExtraKeyDataset(dataset, extra)

Bases: Dataset

Dataset that includes an extra key to add to the data dict. This is useful for adding a REINFORCE baseline reward to the data dict.

class rl4co.data.dataset.TensorDictDataset(data)

Bases: Dataset

Dataset compatible with TensorDicts. It is better to “disassemble” the TensorDict into a list of dicts. See tensordict_collate_fn for more details.

Note

Check out the issue on tensordict for more details: https://github.com/pytorch-labs/tensordict/issues/374. Note that directly indexing TensorDicts may be faster in creating the dataset but uses > 3x more CPU.

rl4co.data.dataset.tensordict_collate_fn(batch)

Collate function compatible with TensorDicts. Reassemble the list of dicts into a TensorDict; seems to be way more efficient than using a TensorDictDataset.

Note

Check out the issue on tensordict for more details: https://github.com/pytorch-labs/tensordict/issues/374. Note that directly indexing TensorDicts may be faster in creating the dataset but uses > 3x more CPU.

Data Generation

rl4co.data.generate_data.generate_dataset(filename=None, data_dir='data', name=None, problem='all', data_distribution='all', dataset_size=10000, graph_sizes=[20, 50, 100], overwrite=False, seed=1234, disable_warning=True)

We keep a similar structure as in Kool et al. 2019 but save and load the data as npz This is way faster and more memory efficient than pickle and also allows for easy transfer to TensorDict

rl4co.data.generate_data.generate_default_datasets(data_dir)

Generate the default datasets used in the paper and save them to data_dir/problem

rl4co.data.generate_data.generate_env_data(env_type, *args, **kwargs)

Generate data for a given environment type in the form of a dictionary

rl4co.data.generate_data.generate_mdpp_data(dataset_size, size=10, num_probes_min=2, num_probes_max=5, num_keepout_min=1, num_keepout_max=50, lock_size=True)

Generate data for the nDPP problem. If lock_size is True, then the size if fixed and we skip the size argument if it is not 10. This is because the RL environment is based on a real-world PCB (parametrized with data)

rl4co.data.generate_data.generate_op_data(dataset_size, op_size, prize_type='const')

rl4co.data.generate_data.generate_pctsp_data(dataset_size, pctsp_size, penalty_factor=3)

rl4co.data.generate_data.generate_tsp_data(dataset_size, tsp_size)

rl4co.data.generate_data.generate_vrp_data(dataset_size, vrp_size, capacities=None)

Transforms

class rl4co.data.transforms.StateAugmentation(env_name=None, num_augment=8, use_dihedral_8=False, normalize=False)

Bases: object

Augment state by N times via symmetric rotation/reflection transform

Parameters:

• env_name (Optional[str]) – environment name

• num_augment (int) – number of augmentations

• use_dihedral_8 (bool) – whether to use dihedral_8_augmentation. If True, then num_augment must be 8

• normalize (bool) – whether to normalize the augmented data

__call__(td)

Call self as a function.

Return type:

TensorDict

rl4co.data.transforms.dihedral_8_augmentation(xy)

Augmentation (x8) for grid-based data (x, y) as done in POMO. This is a Dihedral group of order 8 (rotations and reflections) https://en.wikipedia.org/wiki/Examples_of_groups#dihedral_group_of_order_8

Parameters:

xy (Tensor) – [batch, graph, 2] tensor of x and y coordinates

Return type:

Tensor

rl4co.data.transforms.dihedral_8_augmentation_wrapper(xy, reduce=True, *args, **kw)

Wrapper for dihedral_8_augmentation. If reduce, only return the first 1/8 of the augmented data since the augmentation augments the data 8 times.

Return type:

Tensor

rl4co.data.transforms.env_aug_feats(env_name=None)

What features to augment for a given environment Usually, locs already includes depot, so we don’t need to augment depot

rl4co.data.transforms.min_max_normalize(x)

rl4co.data.transforms.symmetric_augmentation(xy, num_augment=8, first_augment=False)

Augment xy data by num_augment times via symmetric rotation transform and concatenate to original data

Parameters:

• xy (Tensor) – [batch, graph, 2] tensor of x and y coordinates

• num_augment (int) – number of augmentations

• first_augment (bool) – whether to augment the first data point

rl4co.data.transforms.symmetric_transform(x, y, phi, offset=0.5)

SR group transform with rotation and reflection Like the one in SymNCO, but a vectorized version

Parameters:

• x (Tensor) – [batch, graph, 1] tensor of x coordinates

• y (Tensor) – [batch, graph, 1] tensor of y coordinates

• phi (Tensor) – [batch, 1] tensor of random rotation angles

• offset (float) – offset for x and y coordinates

Utils

rl4co.data.utils.check_extension(filename, extension='.npz')

Check that filename has extension, otherwise add it

rl4co.data.utils.load_npz_to_tensordict(filename)

Load a npz file directly into a TensorDict We assume that the npz file contains a dictionary of numpy arrays This is at least an order of magnitude faster than pickle