implement floyd on tmat_class atsp generation

[abcdhhhh]

Description

Change the algorithm of tmat_class ATSP generation. Expect lower time and space complexity when producing the same result.

Motivation and Context

See #225

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• New feature (non-breaking change which adds core functionality)

Checklist

• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[abcdhhhh]

I have also added a val/test dataset generator of ATSP using numpy in rl4co/data/generate_data.py.

[fedebotu]

Awesome job!

from typing import Union, Callable

import torch

from torch.distributions import Uniform
from tensordict.tensordict import TensorDict

from rl4co.utils.pylogger import get_pylogger
from rl4co.envs.common.utils import get_sampler, Generator

log = get_pylogger(__name__)



class ATSPGenerator(Generator):
    """Data generator for the Asymmetric Travelling Salesman Problem (ATSP)
    Generate distance matrices inspired by the reference MatNet (Kwon et al., 2021)
    We satifsy the triangle inequality (TMAT class) in a batch

    Args:
        num_loc: number of locations (customers) in the TSP
        min_dist: minimum value for the distance between nodes
        max_dist: maximum value for the distance between nodes
        dist_distribution: distribution for the distance between nodes
        tmat_class: whether to generate a class of distance matrix

    Returns:
        A TensorDict with the following keys:
            locs [batch_size, num_loc, 2]: locations of each customer
    """
    def __init__(
        self,
        num_loc: int = 10,
        min_dist: float = 0.0,
        max_dist: float = 1.0,
        dist_distribution: Union[
            int, float, str, type, Callable
        ] = Uniform,
        tmat_class: bool = True,
        **kwargs
    ):
        self.num_loc = num_loc
        self.min_dist = min_dist
        self.max_dist = max_dist
        self.tmat_class = tmat_class

        # Distance distribution
        if kwargs.get("dist_sampler", None) is not None:
            self.dist_sampler = kwargs["dist_sampler"]
        else:
            self.dist_sampler = get_sampler("dist", dist_distribution, 0.0, 1.0, **kwargs)

    def _generate(self, batch_size) -> TensorDict:
        # Generate distance matrices inspired by the reference MatNet (Kwon et al., 2021)
        # We satifsy the triangle inequality (TMAT class) in a batch
        batch_size = [batch_size] if isinstance(batch_size, int) else batch_size
        dms = (
            self.dist_sampler.sample((batch_size + [self.num_loc, self.num_loc]))
            * (self.max_dist - self.min_dist)
            + self.min_dist
        )
        dms[..., torch.arange(self.num_loc), torch.arange(self.num_loc)] = 0
        log.info("Using TMAT class (triangle inequality): {}".format(self.tmat_class))
        if self.tmat_class:
            for i in range(self.num_loc):
                dms = torch.minimum(dms, dms[..., :, [i]] + dms[..., [i], :])
        return TensorDict({"cost_matrix": dms}, batch_size=batch_size)
    
class OldATSPGenerator(ATSPGenerator):
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _generate(self, batch_size) -> TensorDict:
        # Generate distance matrices inspired by the reference MatNet (Kwon et al., 2021)
        # We satifsy the triangle inequality (TMAT class) in a batch
        batch_size = [batch_size] if isinstance(batch_size, int) else batch_size
        dms = (
            self.dist_sampler.sample((batch_size + [self.num_loc, self.num_loc]))
            * (self.max_dist - self.min_dist)
            + self.min_dist
        )
        dms[..., torch.arange(self.num_loc), torch.arange(self.num_loc)] = 0
        log.info("Using TMAT class (triangle inequality): {}".format(self.tmat_class))
        if self.tmat_class:
            while True:
                old_dms = dms.clone()
                dms, _ = (
                    dms[..., :, None, :] + dms[..., None, :, :].transpose(-2, -1)
                ).min(dim=-1)
                if (dms == old_dms).all():
                    break
        return TensorDict({"cost_matrix": dms}, batch_size=batch_size)
    
    
generator = ATSPGenerator(num_loc=100)
generator_old = OldATSPGenerator(num_loc=100)

# generate 1000 samples
data = generator(1000)
data_old = generator_old(1000)

print("new")
print(data["cost_matrix"].shape, data["cost_matrix"].mean(), data["cost_matrix"].std())

print("old")
print(data_old["cost_matrix"].shape, data_old["cost_matrix"].mean(), data_old["cost_matrix"].std())

%timeit generator(1000)

685 ms ± 41.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit generator_old(1000)

2.71 s ± 245 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

---

Will merge as soon as the new TorchRL version is released, as the checks will then pass