rl4co 0.0.3.dev0__tar.gz → 0.0.3.dev3__tar.gz

{rl4co-0.0.3.dev0 → rl4co-0.0.3.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rl4co
-Version: 0.0.3.dev0
+Version: 0.0.3.dev3
 Summary: RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark
 Author-email: Federico Berto <berto.federico2@gmail.com>, Chuanbo Hua <cbhua@kaist.ac.kr>, Junyoung Park <junyoungpark.ml@gmail.com>
 License:                                  Apache License
@@ -224,14 +224,15 @@ License-File: LICENSE
 <div align="center">
 
 # RL4CO
-    
-An extensive Reinforcement Learning (RL) for Combinatorial Optimization (CO) benchmark. Our goal is to provide a unified framework for RL-based CO algorithms, and to facilitate reproducible research in this field, decoupling the science from the engineering. 
-        
+
+An extensive Reinforcement Learning (RL) for Combinatorial Optimization (CO) benchmark. Our goal is to provide a unified framework for RL-based CO algorithms, and to facilitate reproducible research in this field, decoupling the science from the engineering.
+
 <a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
 <a href="https://pytorchlightning.ai/"><img alt="Lightning" src="https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white"></a>
 <a href="https://github.com/pytorch/rl"><img alt="base: TorchRL" src="https://img.shields.io/badge/base-TorchRL-red">
 <a href="https://hydra.cc/"><img alt="config: Hydra" src="https://img.shields.io/badge/config-Hydra-89b8cd"></a> [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
-![license](https://img.shields.io/badge/license-Apache%202.0-green.svg?)
+![license](https://img.shields.io/badge/license-Apache%202.0-green.svg?)[![PyPI](https://img.shields.io/pypi/v/rl4co?logo=pypi)](https://pypi.org/project/rl4co)
+[![Test](https://github.com/kaist-silab/rl4co/actions/workflows/tests.yml/badge.svg)](https://github.com/kaist-silab/rl4co/actions/workflows/tests.yml)
 <!-- ![testing](https://github.com/kaist-silab/ncobench/actions/workflows/tests.yml/badge.svg) -->
 
 </div>
@@ -247,6 +248,7 @@ RL4CO is built upon:
 
 
 ## Getting started
+<a href="https://colab.research.google.com/github/kaist-silab/rl4co/blob/main/notebooks/1-quickstart.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
 
 RL4CO is now available for installation on `pip`!
 ```bash
@@ -283,7 +285,7 @@ To get started, we recommend checking out our [quickstart notebook](notebooks/1-
 
 Train model with default configuration (AM on TSP environment):
 ```bash
-python run.py  
+python run.py
 ```
 
 
@@ -293,14 +295,14 @@ python run.py
 
 Train model with chosen experiment configuration from [configs/experiment/](configs/experiment/) (e.g. tsp/am, and environment with 42 cities)
 ```bash
-python run.py experiment=tsp/am env.num_loc=42 
+python run.py experiment=tsp/am env.num_loc=42
 ```
 </details>
 
 
 <details>
     <summary>Disable logging</summary>
-    
+
 ```bash
 python run.py experiment=test/am logger=none '~callbacks.learning_rate_monitor'
 ```
@@ -333,7 +335,7 @@ from rl4co.tasks.rl4co import RL4COLitModule
 config = DictConfig(
     {"data": {
             "train_size": 100000,
-            "val_size": 10000, 
+            "val_size": 10000,
             "batch_size": 512,
         },
     "optimizer": {"lr": 1e-4}}
@@ -368,7 +370,7 @@ pytest tests
 ```
 
 ## Contributing
-Have a suggestion, request, or found a bug? Feel free to [open an issue](https://github.com/kaist-silab/rl4co/issues) or [submit a pull request](https://github.com/kaist-silab/rl4co/pulls). We welcome contributions to RL4CO! 
+Have a suggestion, request, or found a bug? Feel free to [open an issue](https://github.com/kaist-silab/rl4co/issues) or [submit a pull request](https://github.com/kaist-silab/rl4co/pulls). We welcome contributions to RL4CO!
 
 ### Contributors
 <a href="https://github.com/kaist-silab/rl4co/graphs/contributors">

{rl4co-0.0.3.dev0 → rl4co-0.0.3.dev3}/README.md

@@ -1,14 +1,15 @@
 <div align="center">
 
 # RL4CO
-    
-An extensive Reinforcement Learning (RL) for Combinatorial Optimization (CO) benchmark. Our goal is to provide a unified framework for RL-based CO algorithms, and to facilitate reproducible research in this field, decoupling the science from the engineering. 
-        
+
+An extensive Reinforcement Learning (RL) for Combinatorial Optimization (CO) benchmark. Our goal is to provide a unified framework for RL-based CO algorithms, and to facilitate reproducible research in this field, decoupling the science from the engineering.
+
 <a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
 <a href="https://pytorchlightning.ai/"><img alt="Lightning" src="https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white"></a>
 <a href="https://github.com/pytorch/rl"><img alt="base: TorchRL" src="https://img.shields.io/badge/base-TorchRL-red">
 <a href="https://hydra.cc/"><img alt="config: Hydra" src="https://img.shields.io/badge/config-Hydra-89b8cd"></a> [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
-![license](https://img.shields.io/badge/license-Apache%202.0-green.svg?)
+![license](https://img.shields.io/badge/license-Apache%202.0-green.svg?)[![PyPI](https://img.shields.io/pypi/v/rl4co?logo=pypi)](https://pypi.org/project/rl4co)
+[![Test](https://github.com/kaist-silab/rl4co/actions/workflows/tests.yml/badge.svg)](https://github.com/kaist-silab/rl4co/actions/workflows/tests.yml)
 <!-- ![testing](https://github.com/kaist-silab/ncobench/actions/workflows/tests.yml/badge.svg) -->
 
 </div>
@@ -24,6 +25,7 @@ RL4CO is built upon:
 
 
 ## Getting started
+<a href="https://colab.research.google.com/github/kaist-silab/rl4co/blob/main/notebooks/1-quickstart.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
 
 RL4CO is now available for installation on `pip`!
 ```bash
@@ -60,7 +62,7 @@ To get started, we recommend checking out our [quickstart notebook](notebooks/1-
 
 Train model with default configuration (AM on TSP environment):
 ```bash
-python run.py  
+python run.py
 ```
 
 
@@ -70,14 +72,14 @@ python run.py
 
 Train model with chosen experiment configuration from [configs/experiment/](configs/experiment/) (e.g. tsp/am, and environment with 42 cities)
 ```bash
-python run.py experiment=tsp/am env.num_loc=42 
+python run.py experiment=tsp/am env.num_loc=42
 ```
 </details>
 
 
 <details>
     <summary>Disable logging</summary>
-    
+
 ```bash
 python run.py experiment=test/am logger=none '~callbacks.learning_rate_monitor'
 ```
@@ -110,7 +112,7 @@ from rl4co.tasks.rl4co import RL4COLitModule
 config = DictConfig(
     {"data": {
             "train_size": 100000,
-            "val_size": 10000, 
+            "val_size": 10000,
             "batch_size": 512,
         },
     "optimizer": {"lr": 1e-4}}
@@ -145,7 +147,7 @@ pytest tests
 ```
 
 ## Contributing
-Have a suggestion, request, or found a bug? Feel free to [open an issue](https://github.com/kaist-silab/rl4co/issues) or [submit a pull request](https://github.com/kaist-silab/rl4co/pulls). We welcome contributions to RL4CO! 
+Have a suggestion, request, or found a bug? Feel free to [open an issue](https://github.com/kaist-silab/rl4co/issues) or [submit a pull request](https://github.com/kaist-silab/rl4co/pulls). We welcome contributions to RL4CO!
 
 ### Contributors
 <a href="https://github.com/kaist-silab/rl4co/graphs/contributors">

{rl4co-0.0.3.dev0 → rl4co-0.0.3.dev3}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
-include = ['rl4co']
+include = ['rl4co*']
 
 # Automatically determine version from package
 [tool.setuptools.dynamic]

rl4co-0.0.3.dev3/rl4co/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.3.dev3"

rl4co-0.0.3.dev3/rl4co/data/dataset.py

@@ -0,0 +1,73 @@
+from typing import Any
+
+import torch
+
+from tensordict.tensordict import TensorDict
+from torch.utils.data import Dataset
+
+
+class TensorDictDataset(Dataset):
+    """Dataset compatible with TensorDicts
+    For some reason, it is better to "disassemble" the TensorDict into a list of dicts
+    We use a custom collate function to reassemble the TensorDicts
+    NOTE: may want to make an issue on TorchRL to ask best TensorDict practices
+    """
+
+    def __init__(self, data):
+        if isinstance(data, TensorDict):
+            self.data = [
+                {key: value[i] for key, value in data.items()}
+                for i in range(data.shape[0])
+            ]
+        else:
+            self.data = [d for d in data]
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+
+class ExtraKeyDataset(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
+    We had extra_ to identify the key as an extra key
+    """
+
+    def __init__(self, dataset, extra):
+        self.data = dataset.data
+        self.extra = extra
+        assert len(self.data) == len(self.extra), "Data and extra must be same length"
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        data = self.data[idx]
+        data["extra"] = self.extra[idx]
+        return data
+
+
+def 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
+    https://github.com/pytorch-labs/tensordict/issues/374
+    """
+    if isinstance(batch[0], TensorDict):
+        return torch.stack(batch)
+    return TensorDict(
+        {key: torch.stack([b[key] for b in batch]) for key in batch[0].keys()},
+        batch_size=len(batch),
+    )
+
+
+class TensorDictCollate:
+    def __init__(self) -> None:
+        print(
+            "Warning: TensorDictCollateFn is deprecated. Use tensordict_collate_fn instead."
+        )
+        pass
+
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
+        return tensordict_collate_fn(*args, **kwds)

rl4co-0.0.3.dev3/rl4co/data/generate_data.py

@@ -0,0 +1,336 @@
+import argparse
+import logging
+import os
+import sys
+
+import numpy as np
+
+from rl4co.data.utils import check_extension
+from rl4co.utils.pylogger import get_pylogger
+
+log = get_pylogger(__name__)
+
+
+def generate_env_data(env_type, *args, **kwargs):
+    """Generate data for a given environment type in the form of a dictionary"""
+    try:
+        # breakpoint()
+        # remove all None values from args
+        args = [arg for arg in args if arg is not None]
+
+        return getattr(sys.modules[__name__], f"generate_{env_type}_data")(
+            *args, **kwargs
+        )
+    except AttributeError:
+        raise NotImplementedError(f"Environment type {env_type} not implemented")
+
+
+def generate_tsp_data(dataset_size, tsp_size):
+    return {
+        "locs": np.random.uniform(size=(dataset_size, tsp_size, 2)).astype(np.float32)
+    }
+
+
+def generate_vrp_data(dataset_size, vrp_size, capacities=None):
+    # From Kool et al. 2019, Hottung et al. 2022, Kim et al. 2023
+    CAPACITIES = {
+        10: 20.0,
+        15: 25.0,
+        20: 30.0,
+        30: 33.0,
+        40: 37.0,
+        50: 40.0,
+        60: 43.0,
+        75: 45.0,
+        100: 50.0,
+        125: 55.0,
+        150: 60.0,
+        200: 70.0,
+        500: 100.0,
+        1000: 150.0,
+    }
+
+    # If capacities are provided, replace keys in CAPACITIES with provided values if they exist
+    if capacities is not None:
+        for k, v in capacities.items():
+            if k in CAPACITIES:
+                print(f"Replacing capacity for {k} with {v}")
+                CAPACITIES[k] = v
+
+    return {
+        "depot": np.random.uniform(size=(dataset_size, 2)).astype(
+            np.float32
+        ),  # Depot location
+        "locs": np.random.uniform(size=(dataset_size, vrp_size, 2)).astype(
+            np.float32
+        ),  # Node locations
+        "demand": np.random.randint(1, 10, size=(dataset_size, vrp_size)).astype(
+            np.float32
+        ),  # Demand, uniform integer 1 ... 9
+        "capacity": np.full(dataset_size, CAPACITIES[vrp_size]).astype(np.float32),
+    }  # Capacity, same for whole dataset
+
+
+def generate_op_data(dataset_size, op_size, prize_type="const"):
+    depot = np.random.uniform(size=(dataset_size, 2))
+    loc = np.random.uniform(size=(dataset_size, op_size, 2))
+
+    # Methods taken from Fischetti et al. 1998
+    if prize_type == "const":
+        prize = np.ones((dataset_size, op_size))
+    elif prize_type == "unif":
+        prize = (1 + np.random.randint(0, 100, size=(dataset_size, op_size))) / 100.0
+    else:  # Based on distance to depot
+        assert prize_type == "dist"
+        prize_ = np.linalg.norm(depot[:, None, :] - loc, axis=-1)
+        prize = (
+            1 + (prize_ / prize_.max(axis=-1, keepdims=True) * 99).astype(int)
+        ) / 100.0
+
+    # Max length is approximately half of optimal TSP tour, such that half (a bit more) of the nodes can be visited
+    # which is maximally difficult as this has the largest number of possibilities
+    MAX_LENGTHS = {20: 2.0, 50: 3.0, 100: 4.0}
+
+    return {
+        "depot": depot.astype(np.float32),
+        "locs": loc.astype(np.float32),
+        "prize": prize.astype(np.float32),
+        "max_length": np.full(dataset_size, MAX_LENGTHS[op_size]).astype(np.float32),
+    }
+
+
+def generate_pctsp_data(dataset_size, pctsp_size, penalty_factor=3):
+    depot = np.random.uniform(size=(dataset_size, 2))
+    loc = np.random.uniform(size=(dataset_size, pctsp_size, 2))
+
+    # For the penalty to make sense it should be not too large (in which case all nodes will be visited) nor too small
+    # so we want the objective term to be approximately equal to the length of the tour, which we estimate with half
+    # of the nodes by half of the tour length (which is very rough but similar to op)
+    # This means that the sum of penalties for all nodes will be approximately equal to the tour length (on average)
+    # The expected total (uniform) penalty of half of the nodes (since approx half will be visited by the constraint)
+    # is (n / 2) / 2 = n / 4 so divide by this means multiply by 4 / n,
+    # However instead of 4 we use penalty_factor (3 works well) so we can make them larger or smaller
+    MAX_LENGTHS = {20: 2.0, 50: 3.0, 100: 4.0}
+    penalty_max = MAX_LENGTHS[pctsp_size] * (penalty_factor) / float(pctsp_size)
+    penalty = np.random.uniform(size=(dataset_size, pctsp_size)) * penalty_max
+
+    # Take uniform prizes
+    # Now expectation is 0.5 so expected total prize is n / 2, we want to force to visit approximately half of the nodes
+    # so the constraint will be that total prize >= (n / 2) / 2 = n / 4
+    # equivalently, we divide all prizes by n / 4 and the total prize should be >= 1
+    deterministic_prize = (
+        np.random.uniform(size=(dataset_size, pctsp_size)) * 4 / float(pctsp_size)
+    )
+
+    # In the deterministic setting, the stochastic_prize is not used and the deterministic prize is known
+    # In the stochastic setting, the deterministic prize is the expected prize and is known up front but the
+    # stochastic prize is only revealed once the node is visited
+    # Stochastic prize is between (0, 2 * expected_prize) such that E(stochastic prize) = E(deterministic_prize)
+    stochastic_prize = (
+        np.random.uniform(size=(dataset_size, pctsp_size)) * deterministic_prize * 2
+    )
+
+    return {
+        "locs": loc.astype(np.float32),
+        "depot": depot.astype(np.float32),
+        "penalty": penalty.astype(np.float32),
+        "deterministic_prize": deterministic_prize.astype(np.float32),
+        "stochastic_prize": stochastic_prize.astype(np.float32),
+    }
+
+
+def 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)
+    """
+    if lock_size and size != 10:
+        # log.info("Locking size to 10, skipping generate_mdpp_data with size {}".format(size))
+        return None
+
+    bs = dataset_size  # bs = batch_size to generate data in batch
+    m = n = size
+    if isinstance(bs, int):
+        bs = [bs]
+
+    locs = np.stack(np.meshgrid(np.arange(m), np.arange(n)), axis=-1).reshape(-1, 2)
+    locs = locs / np.array([m, n], dtype=np.float32)
+    locs = np.expand_dims(locs, axis=0)
+    locs = np.repeat(locs, bs[0], axis=0)
+
+    available = np.ones((bs[0], m * n), dtype=bool)
+
+    probe = np.random.randint(0, high=m * n, size=(bs[0], 1))
+    np.put_along_axis(available, probe, False, axis=1)
+
+    num_probe = np.random.randint(num_probes_min, num_probes_max + 1, size=(bs[0], 1))
+    probes = np.zeros((bs[0], m * n), dtype=bool)
+    for i in range(bs[0]):
+        p = np.random.choice(m * n, num_probe[i], replace=False)
+        np.put_along_axis(available[i], p, False, axis=0)
+        np.put_along_axis(probes[i], p, True, axis=0)
+
+    num_keepout = np.random.randint(num_keepout_min, num_keepout_max + 1, size=(bs[0], 1))
+    for i in range(bs[0]):
+        k = np.random.choice(m * n, num_keepout[i], replace=False)
+        np.put_along_axis(available[i], k, False, axis=0)
+
+    return {
+        "locs": locs.astype(np.float32),
+        "probe": probes.astype(bool),
+        "action_mask": available.astype(bool),
+    }
+
+
+def 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
+    """
+    assert filename is None or (
+        len(problem) == 1 and len(graph_sizes) == 1
+    ), "Can only specify filename when generating a single dataset"
+
+    distributions_per_problem = {
+        "tsp": [None],
+        "vrp": [None],
+        "pctsp": [None],
+        "op": ["const", "unif", "dist"],
+        "mdpp": [None],
+    }
+    if problem == "all":
+        problems = distributions_per_problem
+    else:
+        problems = {
+            problem: distributions_per_problem[problem]
+            if data_distribution == "all"
+            else [data_distribution]
+        }
+    # breakpoint()
+    fname = filename
+    for problem, distributions in problems.items():
+        for distribution in distributions or [None]:
+            for graph_size in graph_sizes:
+                datadir = os.path.join(data_dir, problem)
+                os.makedirs(datadir, exist_ok=True)
+
+                if filename is None:
+                    fname = os.path.join(
+                        datadir,
+                        "{}{}{}_{}_seed{}.npz".format(
+                            problem,
+                            "_{}".format(distribution)
+                            if distribution is not None
+                            else "",
+                            graph_size,
+                            name,
+                            seed,
+                        ),
+                    )
+                else:
+                    fname = check_extension(filename, extension=".npz")
+
+                if not overwrite and os.path.isfile(
+                    check_extension(fname, extension=".npz")
+                ):
+                    if not disable_warning:
+                        log.info(
+                            "File {} already exists! Run with -f option to overwrite. Skipping...".format(
+                                fname
+                            )
+                        )
+                    continue
+
+                # Set seed
+                np.random.seed(seed)
+
+                # Automatically generate dataset
+                dataset = generate_env_data(
+                    problem, dataset_size, graph_size, distribution
+                )
+
+                # A function can return None in case of an error or a skip
+                if dataset is not None:
+                    # Save to disk as dict
+                    log.info("Saving {} dataset to {}".format(problem, fname))
+                    np.savez(fname, **dataset)
+
+
+def generate_default_datasets(data_dir):
+    """Generate the default datasets used in the paper and save them to data_dir/problem"""
+    generate_dataset(data_dir=data_dir, name="val", problem="all", seed=4321)
+    generate_dataset(data_dir=data_dir, name="test", problem="all", seed=1234)
+    generate_dataset(
+        data_dir=data_dir,
+        name="test",
+        problem="mdpp",
+        seed=1234,
+        graph_sizes=[10],
+        dataset_size=100,
+    )  # EDA (mDPP)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--filename", help="Filename of the dataset to create (ignores datadir)"
+    )
+    parser.add_argument(
+        "--data_dir",
+        default="data",
+        help="Create datasets in data_dir/problem (default 'data')",
+    )
+    parser.add_argument(
+        "--name", type=str, required=True, help="Name to identify dataset"
+    )
+    parser.add_argument(
+        "--problem",
+        type=str,
+        default="all",
+        help="Problem, 'tsp', 'vrp', 'pctsp' or 'op_const', 'op_unif' or 'op_dist'"
+        " or 'all' to generate all",
+    )
+    parser.add_argument(
+        "--data_distribution",
+        type=str,
+        default="all",
+        help="Distributions to generate for problem, default 'all'.",
+    )
+    parser.add_argument(
+        "--dataset_size", type=int, default=10000, help="Size of the dataset"
+    )
+    parser.add_argument(
+        "--graph_sizes",
+        type=int,
+        nargs="+",
+        default=[20, 50, 100],
+        help="Sizes of problem instances (default 20, 50, 100)",
+    )
+    parser.add_argument("-f", action="store_true", help="Set true to overwrite")
+    parser.add_argument("--seed", type=int, default=1234, help="Random seed")
+    parser.add_argument("--disable_warning", action="store_true", help="Disable warning")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+
+    args.overwrite = args.f
+    delattr(args, "f")
+    generate_dataset(**vars(args))

rl4co-0.0.3.dev3/rl4co/data/utils.py

@@ -0,0 +1,23 @@
+import os
+
+import numpy as np
+
+from tensordict.tensordict import TensorDict
+
+
+def 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
+    """
+    x = np.load(filename)
+    x_dict = dict(x)
+    batch_size = x_dict[list(x_dict.keys())[0]].shape[0]
+    return TensorDict(x_dict, batch_size=batch_size)
+
+
+def check_extension(filename, extension=".npz"):
+    """Check that filename has extension, otherwise add it"""
+    if os.path.splitext(filename)[1] != extension:
+        return filename + extension
+    return filename

rl4co-0.0.3.dev3/rl4co/envs/__init__.py

@@ -0,0 +1,11 @@
+from rl4co.envs.atsp import ATSPEnv
+from rl4co.envs.base import RL4COEnvBase
+from rl4co.envs.cvrp import CVRPEnv
+from rl4co.envs.dpp import DPPEnv
+from rl4co.envs.mdpp import MDPPEnv
+from rl4co.envs.mtsp import MTSPEnv
+from rl4co.envs.op import OPEnv
+from rl4co.envs.pctsp import PCTSPEnv
+from rl4co.envs.pdp import PDPEnv
+from rl4co.envs.sdvrp import SDVRPEnv
+from rl4co.envs.tsp import TSPEnv