add greedrl

.gitignore

@@ -0,0 +1,20 @@
+.idea
+*.tar.gz
+logs
+**/__pycache__
+data
+*.log
+*.pkl
+*.pt
+**/build/
+**/dist/
+**/*.egg-info
+.DS_Store
+.nfs*
+*.so
+*.dylib
+*.iml
+target
+**/nohup.out
+*.pth
+**/.flattened-pom.xml

CMakeLists.txt

@@ -0,0 +1,68 @@
+cmake_minimum_required(VERSION 2.8.12)
+project(greedrl_C_ LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 14)
+
+find_package(PythonInterp REQUIRED)
+execute_process(COMMAND "python" "-c" 
+"
+import os
+import torch
+from distutils import sysconfig as s
+print(s.get_python_inc(plat_specific=True))
+print(s.get_config_var('EXT_SUFFIX'))
+print(os.path.dirname(torch.__file__))
+"
+RESULT_VARIABLE _PYTHON_SUCCESS
+OUTPUT_VARIABLE _PYTHON_VALUES
+ERROR_VARIABLE _PYTHON_ERROR_VALUE)
+
+if(NOT _PYTHON_SUCCESS MATCHES 0)
+    message("_PYTHON_SUCCESS: ${_PYTHON_SUCCESS}")
+    message("_PYTHON_VALUES: ${_PYTHON_VALUES}")
+    message("_PYTHON_ERROR_VALUE: ${_PYTHON_ERROR_VALUE}")
+    message(FATAL_ERROR "get python config error!")
+endif()
+
+string(REGEX REPLACE "\n" ";" _PYTHON_VALUES ${_PYTHON_VALUES})
+list(GET _PYTHON_VALUES 0 PYTHON_INCLUDE_DIR)
+list(GET _PYTHON_VALUES 1 PYTHON_EXT_SUFFIX)
+list(GET _PYTHON_VALUES 2 TORCH_HOME)
+
+include_directories(
+    ${PYTHON_INCLUDE_DIR}
+    ${TORCH_HOME}/include
+    ${TORCH_HOME}/include/TH
+    ${TORCH_HOME}/include/THC
+    ${TORCH_HOME}/include/torch/csrc/api/include
+)
+
+string(LENGTH "${CMAKE_SOURCE_DIR}/" SOURCE_PATH_LENGTH)
+add_compile_options(-DSOURCE_PATH_LENGTH=${SOURCE_PATH_LENGTH})
+add_compile_options(-D_GLIBCXX_USE_CXX11_ABI=0 -fvisibility=hidden -fopenmp)
+
+if(${CMAKE_CXX_COMPILER_ID} STREQUAL "Clang")
+    add_link_options(-undefined dynamic_lookup)
+endif()
+
+file(GLOB_RECURSE CSRC_CPP csrc/*.cpp)
+
+add_library(greedrl_c MODULE ${CSRC_CPP})
+set_target_properties(greedrl_c PROPERTIES PREFIX "")
+set_target_properties(greedrl_c PROPERTIES SUFFIX "${PYTHON_EXT_SUFFIX}")
+target_compile_options(greedrl_c PRIVATE -Wno-sign-conversion -O3)
+target_link_libraries(greedrl_c c10 torch torch_cpu torch_python)
+target_link_directories(greedrl_c PRIVATE ${TORCH_HOME}/lib)
+
+find_package(CUDA)
+if(CUDA_FOUND)
+    enable_language(CUDA)
+    file(GLOB_RECURSE CSRC_CU csrc/*.cu)
+    add_library(greedrl_cu OBJECT ${CSRC_CU})
+    target_compile_options(greedrl_cu PRIVATE -keep -Xptxas -v --expt-relaxed-constexpr --expt-extended-lambda -O3)
+    set_target_properties(greedrl_cu PROPERTIES POSITION_INDEPENDENT_CODE ON CUDA_ARCHITECTURES "70;75;80")
+    add_compile_definitions(CUDA_FOUND)
+    include_directories(${CUDA_INCLUDE_DIRS})
+    target_link_libraries(greedrl_c torch_cuda greedrl_cu)
+    target_link_directories(greedrl_c PRIVATE ${TORCH_HOME}/lib)
+endif()

README.md

@@ -1,3 +1,628 @@
 ---
 license: apache-2.0
+pipeline_tag: reinforcement-learning
+tags:
+- Deep Reinforcement Learning
+- Combinatorial Optimization
+- Reinforcement Learning
+- Vehicle Routing Problem
 ---
+
+  ![](./images/GREEDRL-Logo-Original-640.png)
+
+
+# 🤠GreedRL
+
+## Overview
+
+- 🤠GreedRL is a fast and general framework for **Combinatorial Optimization Problems (COPs)**, based on **Deep Reinforcement Learning (DRL)**.
+
+- 🤠GreedRL achieves **1200 times faster and 3% improved performance** than [Google OR-Tools](https://developers.google.com/optimization) for large-scale (>=1000 nodes) CVRPs.
+
+## 🏆Award
+
+  [INFORMS 2021 Franz Edelman Award finalists](https://www.informs.org/Resource-Center/Video-Library/Edelman-Competition-Videos/2021-Edelman-Competition-Videos/2021-Edelman-Finalist-Alibaba) for Achievement in Operations Research and the Management Sciences (recognized for our work on Cainiao Network VRP algorithm).
+
+
+## Main features
+
+* **GENERAL**
+  
+  🤠GreedRL makes **a high level of abstraction for COPs**, which can solve various types of problems, such as TSP, CVRP, VRPTW, PDPTW, SDVRP, DPDP, Order Batching, etc.
+
+* **HIGH-PERFORMANCE**
+
+  🤠GreedRL have improved the DRL environment (Env) simulation speed by **CUDA and C++ implementations**.
+  
+* **USER-FRIENDLY**
+
+  🤠GreedRL framework provides **user-friendly ability for COPs modeling**, where users only need to declare constraints, objectives and variables of COPs. For more examples, please refer to [COPs Modeling examples](https://huggingface.co/Cainiao-AI/GreedRL/blob/main/README.md#cops-modeling-examples).
+
+## Editions
+
+We provide an open source Community Edition and an Enterprise Edition of our 🤠GreedRL for users.
+
+- **The Community Edition** is now released and available to [download](https://huggingface.co/Cainiao-AI/GreedRL).
+- **The Enterprise Edition** has a high-performance implementation that achives a faster computing speed, especially when solving larg-scale COPs. For more informations, please contact <a href="mailto:[email protected]">us</a>.
+
+
+## Architecture
+  ![](./images/GREEDRL-Framwork_en.png)
+
+## COPs Modeling examples
+
+
+### Capacitated Vehicle Routing Problem (CVRP)
+<details>
+    <summary>CVRP</summary>
+
+```python
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl.function import *
+from greedrl import Problem, Solution, Solver
+from greedrl import runner
+
+features = [continuous_feature('task_demand'),
+            continuous_feature('worker_weight_limit'),
+            continuous_feature('distance_matrix'),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             worker_variable('worker_weight_limit'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand_now <= 0
+        # 车辆容量限制
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_demand_now * self.task_weight > worker_weight_limit[:, None]
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        return self.distance_last_to_this
+
+    def step_task(self):
+        return self.distance_last_to_this
+```
+
+</details>
+
+### Pickup and Delivery Problem with Time Windows (PDPTW)
+<details>
+    <summary>PDPTW</summary>
+
+```python
+from greedrl.model import runner
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl.function import *
+from greedrl import Problem, Solution, Solver
+
+features = [local_category('task_group'),
+            global_category('task_priority', 2),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             feature_variable('task_group'),
+             feature_variable('task_priority'),
+             feature_variable('task_due_time2', feature='task_due_time'),
+             task_variable('task_due_time'),
+             task_variable('task_service_time'),
+             task_variable('task_due_time_penalty'),
+             worker_variable('worker_basic_cost'),
+             worker_variable('worker_distance_cost'),
+             worker_variable('worker_due_time'),
+             worker_variable('worker_weight_limit'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             worker_used_resource('worker_used_time', 'distance_matrix', 'task_service_time', 'task_ready_time',
+                                  'worker_ready_time'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_worker_end(self):
+        return task_group_split(self.task_group, self.task_demand_now <= 0)
+
+    def mask_task(self):
+        mask = self.task_demand_now <= 0
+        mask |= task_group_priority(self.task_group, self.task_priority, mask)
+
+        worker_used_time = self.worker_used_time[:, None] + self.distance_this_to_task
+        mask |= (worker_used_time > self.task_due_time2) & (self.task_priority == 0)
+
+        # 容量约束
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_demand_now * self.task_weight > worker_weight_limit[:, None]
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_start(self):
+        return self.worker_basic_cost
+
+    def step_worker_end(self):
+        feasible = self.worker_used_time <= self.worker_due_time
+        return self.distance_last_to_this * self.worker_distance_cost, feasible
+
+    def step_task(self):
+        worker_used_time = self.worker_used_time - self.task_service_time
+        feasible = worker_used_time <= self.task_due_time
+        feasible &= worker_used_time <= self.worker_due_time
+        cost = self.distance_last_to_this * self.worker_distance_cost
+        return torch.where(feasible, cost, cost + self.task_due_time_penalty), feasible
+```
+
+</details>
+
+
+### VRP with Time Windows (VRPTW)
+<details>
+    <summary>VRPTW</summary>
+
+```python
+from greedrl import Problem, Solution, Solver
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl.function import *
+from greedrl.model import runner
+from greedrl.myenv import VrptwEnv
+
+features = [continuous_feature('worker_weight_limit'),
+            continuous_feature('worker_ready_time'),
+            continuous_feature('worker_due_time'),
+            continuous_feature('worker_basic_cost'),
+            continuous_feature('worker_distance_cost'),
+            continuous_feature('task_demand'),
+            continuous_feature('task_weight'),
+            continuous_feature('task_ready_time'),
+            continuous_feature('task_due_time'),
+            continuous_feature('task_service_time'),
+            continuous_feature('distance_matrix')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             feature_variable('task_due_time'),
+             feature_variable('task_ready_time'),
+             feature_variable('task_service_time'),
+             worker_variable('worker_weight_limit'),
+             worker_variable('worker_due_time'),
+             worker_variable('worker_basic_cost'),
+             worker_variable('worker_distance_cost'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             worker_used_resource('worker_used_time', 'distance_matrix', 'task_service_time', 'task_ready_time',
+                'worker_ready_time'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand_now <= 0
+        # 车辆容量限制
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_demand_now * self.task_weight > worker_weight_limit[:, None]
+
+        worker_used_time = self.worker_used_time[:, None] + self.distance_this_to_task
+        mask |= worker_used_time > self.task_due_time
+
+        worker_used_time = torch.max(worker_used_time, self.task_ready_time)
+        worker_used_time += self.task_service_time
+        worker_used_time += self.distance_task_to_end
+        mask |= worker_used_time > self.worker_due_time[:, None]
+
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_start(self):
+        return self.worker_basic_cost
+
+    def step_worker_end(self):
+        return self.distance_last_to_this * self.worker_distance_cost
+
+    def step_task(self):
+        return self.distance_last_to_this * self.worker_distance_cost
+```
+
+</details>
+
+### Travelling Salesman Problem (TSP)
+<details>
+    <summary>TSP</summary>
+
+```python
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl import Problem
+from greedrl import runner
+
+features = [continuous_feature('task_location'),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True),
+             edge_variable('distance_last_to_loop', feature='distance_matrix', last_to_loop=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_task(self):
+        mask = self.task_demand_now <= 0
+        return mask
+
+    def mask_worker_end(self):
+        return torch.any(self.task_demand_now > 0, 1)
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        return self.distance_last_to_loop
+
+    def step_task(self):
+        return self.distance_last_to_this
+```
+
+</details>
+
+### Split Delivery Vehicle Routing Problem (SDVRP)
+<details>
+    <summary>SDVRP</summary>
+
+```python
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl import Problem
+from greedrl import runner
+
+features = [continuous_feature('task_demand'),
+            continuous_feature('worker_weight_limit'),
+            continuous_feature('distance_matrix'),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             task_variable('task_weight_this', feature='task_weight'),
+             worker_variable('worker_weight_limit'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        return torch.min(self.task_demand_this, worker_weight_limit // self.task_weight_this)
+
+    def mask_task(self):
+        mask = self.task_demand <= 0
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_weight > worker_weight_limit[:, None]
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        return self.distance_last_to_this
+
+    def step_task(self):
+        return self.distance_last_to_this
+```
+
+</details>
+
+### Realistic Business Scenario
+<details>
+    <summary>real-time Dynamic Pickup and Delivery Problem (DPDP)</summary>
+
+```python
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl.function import *
+from greedrl import Problem
+from greedrl import runner
+
+features = [local_category('task_order'),
+            global_category('task_type', 2),
+            global_category('task_new_order', 2),
+            variable_feature('time_this_to_task'),
+            continuous_feature('x_time_matrix'),
+            continuous_feature('task_due_time_x'),
+            continuous_feature('worker_task_mask')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             task_variable('task_pickup_this', feature='task_pickup'),
+             task_variable('task_due_time_this', feature='task_due_time'),
+             feature_variable('task_order', feature='task_order'),
+             feature_variable('task_type', feature='task_type'),
+             feature_variable('task_new_pickup', feature='task_new_pickup'),
+             feature_variable('worker_task_mask', feature='worker_task_mask'),
+             worker_count_now('worker_count_now', feature='worker_count'),
+             worker_variable('worker_min_old_task_this', feature='worker_min_old_task'),
+             worker_variable('worker_max_new_order_this', feature='worker_max_new_order'),
+             worker_variable('worker_task_mask_this', feature='worker_task_mask'),
+             worker_used_resource('worker_used_old_task', task_require='task_old'),
+             worker_used_resource('worker_used_new_order', task_require='task_new_pickup'),
+             worker_used_resource('worker_used_time', edge_require='time_matrix'),
+             edge_variable('time_this_to_task', feature='x_time_matrix', this_to_task=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_worker_start(self):
+        mask = self.worker_count_now <= 0
+
+        finished = self.task_demand_now <= 0
+        worker_task_mask = self.worker_task_mask | finished[:, None, :]
+        mask |= torch.all(worker_task_mask, 2)
+
+        return mask
+
+    def mask_worker_end(self):
+        mask = self.worker_used_old_task < self.worker_min_old_task_this
+        mask |= task_group_split(self.task_order, self.task_demand_now <= 0)
+        return mask
+
+    def mask_task(self):
+        mask = self.task_demand_now <= 0
+
+        mask |= task_group_priority(self.task_order, self.task_type, mask)
+
+        worker_max_new_order = self.worker_max_new_order_this - self.worker_used_new_order
+        mask |= self.task_new_pickup > worker_max_new_order[:, None]
+
+        mask |= self.worker_task_mask_this
+
+        return mask
+
+    def finished(self):
+        worker_mask = self.worker_count_now <= 0
+        task_mask = self.task_demand_now <= 0
+        worker_task_mask = worker_mask[:, :, None] | task_mask[:, None, :]
+
+        worker_task_mask |= self.worker_task_mask
+        batch_size = worker_task_mask.size(0)
+        worker_task_mask = worker_task_mask.view(batch_size, -1)
+        return worker_task_mask.all(1)
+
+
+class Objective:
+
+    def step_task(self):
+        over_time = (self.worker_used_time - self.task_due_time_this).clamp(min=0)
+        pickup_time = self.worker_used_time * self.task_pickup_this
+        return self.worker_used_time + over_time + pickup_time
+
+    def step_finish(self):
+        return self.task_demand_now.sum(1) * 1000
+```
+
+</details>
+
+### Order Batching Problem
+<details>
+    <summary>Batching</summary>
+
+```python
+from greedrl import Problem, Solver
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl import runner
+
+
+features = [local_feature('task_area'),
+            local_feature('task_roadway'),
+            local_feature('task_area_group'),
+            sparse_local_feature('task_item_id', 'task_item_num'),
+            sparse_local_feature('task_item_owner_id', 'task_item_num'),
+            variable_feature('worker_task_item'),
+            variable_feature('worker_used_roadway'),
+            variable_feature('worker_used_area')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_item_id'),
+             feature_variable('task_item_num'),
+             feature_variable('task_item_owner_id'),
+             feature_variable('task_area'),
+             feature_variable('task_area_group'),
+             feature_variable('task_load'),
+             feature_variable('task_group'),
+             worker_variable('worker_load_limit'),
+             worker_variable('worker_area_limit'),
+             worker_variable('worker_area_group_limit'),
+             worker_task_item('worker_task_item', item_id='task_item_id', item_num='task_item_num'),
+             worker_task_item('worker_task_item_owner', item_id='task_item_owner_id', item_num='task_item_num'),
+             worker_used_resource('worker_used_load', task_require='task_load'),
+             worker_used_resource('worker_used_area', task_require='task_area'),
+             worker_used_resource('worker_used_roadway', task_require='task_roadway'),
+             worker_used_resource('worker_used_area_group', task_require='task_area_group')]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_worker_end(self):
+        return self.worker_used_load < self.worker_load_limit
+
+    def mask_task(self):
+        # completed tasks
+        mask = self.task_demand_now <= 0
+        # mask |= task_group_priority(self.task_group, self.task_out_stock_time, mask)
+
+        NT = self.task_item_id.size(1)
+        worker_task_item = self.worker_task_item[:, None, :]
+        worker_task_item = worker_task_item.expand(-1, NT, -1)
+        task_item_in_worker = worker_task_item.gather(2, self.task_item_id.long())
+        task_item_in_worker = (task_item_in_worker > 0) & (self.task_item_num > 0)
+
+        worker_task_item_owner = self.worker_task_item_owner[:, None, :]
+        worker_task_item_owner = worker_task_item_owner.expand(-1, NT, -1)
+        task_item_owner_in_worker = worker_task_item_owner.gather(2, self.task_item_owner_id.long())
+        task_item_owner_in_worker = (task_item_owner_in_worker > 0) & (self.task_item_num > 0)
+
+        # 
+        mask |= torch.any(task_item_in_worker & ~task_item_owner_in_worker, 2)
+
+        worker_load_limit = self.worker_load_limit - self.worker_used_load
+        mask |= (self.task_load > worker_load_limit[:, None])
+
+        task_area = self.task_area + self.worker_used_area[:, None, :]
+        task_area_num = task_area.clamp(0, 1).sum(2, dtype=torch.int32)
+        mask |= (task_area_num > self.worker_area_limit[:, None])
+
+        tak_area_group = self.task_area_group + self.worker_used_area_group[:, None, :]
+        tak_area_group_num = tak_area_group.clamp(0, 1).sum(2, dtype=torch.int32)
+        mask |= (tak_area_group_num > self.worker_area_group_limit[:, None])
+
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        area_num = self.worker_used_area.clamp(0, 1).sum(1)
+        roadway_num = self.worker_used_roadway.clamp(0, 1).sum(1)
+        item_num = self.worker_task_item.clamp(0, 1).sum(1)
+        penalty = (self.worker_load_limit - self.worker_used_load) * 10
+        return area_num * 100 + roadway_num * 10 + item_num + penalty
+```
+
+</details>
+
+
+#
+#
+# Getting started
+
+## Description
+We are delighted to release 🤠GreedRL Community Edition, as well as example of training and testing scripts for the standard Capacitated VRP (CVRP), you can download it and get started.
+
+## Test environment
+🤠GreedRL Community Edition has been tested on Ubuntu 18.04 with GCC compiler v7.5.0 and CUDA version 11.4, and a [Miniconda](https://docs.conda.io/en/latest/miniconda.html#system-requirements) distribution with Python 3.8. We recommend using a similar configuration to avoid any possiblem compilation issue. 
+
+## Installation
+First, clone the repository.
+```aidl
+$ git clone https://huggingface.co/Cainiao-AI/GreedRL
+```
+Then, create and activate a python environment using conda, and install required packages.
+```aidl
+$ conda create -n python38 python==3.8
+$ source activate python38
+$ pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu113
+```
+Finally, compile and add the resulting library `greedrl` to the `PYTHONPATH`
+```aidl
+$ python setup.py build
+$ export PYTHONPATH={your_current_path}/build/lib.linux-x86_64-cpython-38/:$PYTHONPATH
+```
+
+## CVRP Training
+
+1. Training data
+
+We use generated data for the training phase, the customers and depot locations are randomly generated in the unit square [0,1] X [0,1]. For CVRP, we assume that the demand of each node is a discrete number in {1,...,9}, chosen uniformly at random, and each vehicle has a default capacity of 50.
+
+
+2. Start training
+```python
+$ cd examples/cvrp
+$ python train.py --model_filename cvrp_100.pt --problem_size 100
+```
+
+## CVRP Testing
+
+After training process, you'll get a trained model, like `cvrp_100.pt`, that you can use for test.
+
+```python
+$ cd examples/cvrp
+$ python solve.py --device cpu --model_name cvrp_100.pt --problem_size 100
+```
+
+# Support
+We look forward you to downloading it, using it, and opening discussion if you encounter any problems or have ideas on building an even better experience.
+For commercial enquiries, please contact <a href="mailto:[email protected]">us</a>.
+
+# Citation
+```
+@article{hu2022alibaba,
+  title={Alibaba vehicle routing algorithms enable rapid pick and delivery},
+  author={Hu, Haoyuan and Zhang, Ying and Wei, Jiangwen and Zhan, Yang and Zhang, Xinhui and Huang, Shaojian and Ma, Guangrui and Deng, Yuming and Jiang, Siwei},
+  journal={INFORMS Journal on Applied Analytics},
+  volume={52},
+  number={1},
+  pages={27--41},
+  year={2022},
+  publisher={INFORMS}
+}
+```

csrc/common.h

@@ -0,0 +1,184 @@
+#pragma once
+#include <cfloat>
+#include <climits>
+#include <cstdint>
+#include <limits>
+#include <chrono>
+#include <stdexcept>
+#include <torch/extension.h>
+
+#define ASSERT(c) assert(c)
+#define ALIGN(v, n) ((v + n - 1) / n * n)
+#define INF std::numeric_limits<float>::infinity()
+#define __FILENAME__ (__FILE__+ SOURCE_PATH_LENGTH)
+
+#define GRL_ERROR(format, args...)                                      \
+    greedrl_error(__FILENAME__, __LINE__, format, ##args);              \
+
+
+#define GRL_CHECK(flag, format, args...)                                \
+    greedrl_check(__FILENAME__, __LINE__, flag, format, ##args);        \
+
+
+#define MALLOC(ptr, T, size)                                            \
+    ptr = (T*) malloc(sizeof(T) * (size));                              \
+    GRL_CHECK(ptr != nullptr, "out of memory!");                        \
+
+
+#define GALLOC(ptr, T, size)                                            \
+    GRL_CHECK((size) > 0, "malloc 0 bytes");                            \
+    T* const ptr = (T*) malloc(sizeof(T) * (size));                     \
+    GRL_CHECK(ptr != nullptr, "out of memory!");                        \
+    AllocGuard ptr##_##alloc##_##guard(ptr);                            \
+
+
+#define REALLOC(ptr, T, size)                                           \
+    GRL_CHECK((size) > 0, "malloc 0 bytes");                            \
+    ptr = (T*) realloc(ptr, sizeof(T) * (size));                        \
+    GRL_CHECK(ptr != nullptr, "out of memory!");                        \
+
+
+#define GRL_CHECK_TENSOR(tensor, device, allow_sub_contiguous, allow_null, ...)     \
+    greedrl_check_tensor(__FILENAME__, __LINE__, tensor, #tensor, device,           \
+                         allow_sub_contiguous, allow_null, {__VA_ARGS__});          \
+
+
+const int GRL_WORKER_START = 0;
+const int GRL_WORKER_END = 1;
+const int GRL_TASK = 2;
+const int GRL_FINISH = 3;
+
+const int MAX_BATCH_SIZE = 100000; 
+const int MAX_TASK_COUNT = 5120;
+const int MAX_SHARED_MEM = 48128;
+
+using String = std::string;
+using Device = torch::Device;
+using Tensor = torch::Tensor;
+using TensorMap = std::map<String, Tensor>;
+using TensorList = std::vector<Tensor>;
+
+
+inline void greedrl_error(const char* const file, const int64_t line,
+                          const char* const format, ...)
+{
+    const int N = 2048;
+    static char buf[N];
+
+    va_list args;
+    va_start(args, format);
+    int n = vsnprintf(buf, N, format, args);
+    va_end(args);
+
+    if(n < N)
+    {
+        snprintf(buf+n, N-n, " at %s:%ld", file, line);
+    }
+
+    throw std::runtime_error(buf);
+}
+
+inline void greedrl_check(const char* const file, const int64_t line,
+                          const bool flag, const char* const format, ...)
+{
+    if(flag)
+    {
+        return;
+    }
+    
+    const int N = 2048;
+    static char buf[N];
+
+    va_list args;
+    va_start(args, format);
+    int n = vsnprintf(buf, N, format, args);
+    va_end(args);
+
+    if(n < N)
+    {
+        snprintf(buf+n, N-n, " at %s:%ld", file, line);
+    }
+
+    throw std::runtime_error(buf);
+}
+
+// contiguous except the 1st dimension
+inline bool is_sub_contiguous(const Tensor& tensor)
+{
+    int dim = tensor.dim();
+    if(dim==1) return true;
+
+    auto sizes = tensor.sizes();
+    auto strides = tensor.strides();  
+    
+    if(strides[dim-1] != 1) return false;
+ 
+    int s = 1;
+    for(int i=dim-2; i>0; i--)
+    {
+        s *= sizes[i+1];
+        if(strides[i] != s) return false;
+    }
+    
+    return true;
+
+};
+
+inline void greedrl_check_tensor(const char* const file, 
+                                 const int line,
+                                 const Tensor& tensor,
+                                 const String& name, 
+                                 const Device& device, 
+                                 bool allow_sub_contiguous,
+                                 bool allow_null,
+                                 std::initializer_list<int> sizes)
+{
+    greedrl_check(file, line, tensor.numel() < 1000 * 1000 * 1000, "tensor size too large");
+
+    auto device2 = tensor.device();
+    greedrl_check(file, line, device2==device,
+            "'%s' device is %s, but expect %s",
+            name.c_str(), device2.str().c_str(), device.str().c_str());
+    
+    bool is_contiguous = allow_sub_contiguous ? is_sub_contiguous(tensor) : tensor.is_contiguous();
+    greedrl_check(file, line, is_contiguous, "'%s' is not contiguous", name.c_str());
+    
+    if(allow_null && tensor.data_ptr() == nullptr) return;
+    
+    if(tensor.dim() != sizes.size())
+    {
+        greedrl_error(file, line, "'%s' dim is %d, but expect %d", name.c_str(), (int)tensor.dim(), (int)sizes.size());
+    }
+    int i=0;
+    for(auto s:sizes)
+    {
+        greedrl_check(file, line, tensor.size(i)==s, "'%s' size(%d) is %d, but expect %d", name.c_str(), i, (int)tensor.size(i), s);
+        i++;
+    }
+}
+
+
+#ifdef CUDA_FOUND
+
+#include <cuda_runtime_api.h>
+
+#define GRL_CHECK_CUDA(error)\
+    greedrl_check_cuda(error, __FILENAME__, __LINE__);                                                  
+
+inline void greedrl_check_cuda(const cudaError_t& error,
+                               const char* file, const int64_t line)
+{
+    if(error==cudaSuccess)
+    {
+        return;
+    }
+    
+    const int N = 2048;
+    static char buf[N];
+    snprintf(buf, N, "%s, at %s:%ld", cudaGetErrorString(error), file, line);
+    throw std::runtime_error(buf);
+}
+
+cudaDeviceProp& cuda_get_device_prop(int i);
+
+#endif

csrc/pybind.cpp

@@ -0,0 +1,11 @@
+#include <pybind11/pybind11.h>
+#include "task_group_split.h"
+#include "task_group_priority.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(greedrl_c, m) {
+    m.def("task_group_split", &task_group_split);
+    m.def("task_group_priority", &task_group_priority);
+}
+

csrc/task_group_priority.cpp

@@ -0,0 +1,75 @@
+#include "task_group_priority.h"
+
+void task_group_priority_cpu(
+        int* group, int* priority, bool* value, bool* output,
+        int batch_size, int task_num, int group_num)
+{
+    auto temp = torch::make_unique<int[]>(group_num);
+    for(int b=0; b<batch_size; b++)
+    {
+        for(int i=0; i<group_num; i++){
+            temp[i] = std::numeric_limits<int>::max();
+        }
+        
+        for(int i=0; i<task_num; i++){
+            if(value[i]){
+                continue;
+            }
+            int g = group[i];
+            int p = priority[i];
+            if(p < temp[g]){
+                temp[g] = p;
+            }
+        }
+    
+        for(int i=0; i<task_num; i++){
+            int g = group[i];
+            output[i] = priority[i]!=temp[g];
+        }
+        
+        group += task_num;
+        priority += task_num;
+        value += task_num;
+        output += task_num;
+    }
+};
+
+auto task_group_priority(
+        const torch::Tensor& group,
+        const torch::Tensor& priority,
+        const torch::Tensor& value) -> torch::Tensor 
+{
+    auto device = group.device();
+
+    const int batch_size = group.size(0);
+    const int task_num = group.size(1);
+    const int group_num = group.max().item<int>() + 1;
+
+    const int _group_num = group.min().item<int>();
+
+    GRL_CHECK(group_num <= task_num && _group_num >= 0, "group value error");
+
+    GRL_CHECK_TENSOR(group, device, false, false, batch_size, task_num);
+    GRL_CHECK_TENSOR(priority, device, false, false, batch_size, task_num);
+    GRL_CHECK_TENSOR(value, device, false, false, batch_size, task_num);
+
+    auto output = torch::zeros({batch_size, task_num}, torch::dtype(torch::kBool).device(device));
+
+    switch(device.type())
+    {
+        case torch::kCPU:
+            task_group_priority_cpu(group.data_ptr<int>(), priority.data_ptr<int>(), value.data_ptr<bool>(),
+                                    output.data_ptr<bool>(), batch_size, task_num, group_num); 
+            break;
+#ifdef CUDA_FOUND
+        case torch::kCUDA:
+            task_group_priority_cuda(group.data_ptr<int>(), priority.data_ptr<int>(), value.data_ptr<bool>(),
+                                    output.data_ptr<bool>(), batch_size, task_num, group_num, device.index());
+            break;
+#endif
+        default:
+            GRL_ERROR("unsupported device: %s", device.str().c_str());
+    }
+    
+    return output;
+};

csrc/task_group_priority.cu

@@ -0,0 +1,93 @@
+#include "task_group_priority.h"
+
+__global__ void task_group_priority_kernel(
+        int* group, int* priority, bool* value, bool* output,
+        int batch_size, int task_num, int group_num)
+{
+    group += blockIdx.x * task_num;
+    priority += blockIdx.x * task_num;
+    value += blockIdx.x * task_num;
+    output += blockIdx.x * task_num;
+
+    extern __shared__ int temp[];
+    
+    for(int i=threadIdx.x; i<group_num; i+=blockDim.x)
+    {
+        temp[i] = std::numeric_limits<int>::max();
+    }
+    
+    __syncthreads();
+
+    for(int i=threadIdx.x; i<task_num; i+=blockDim.x){
+        if(value[i]){
+            continue;
+        }
+        int g = group[i];
+        int p = priority[i];
+        atomicMin(&temp[g], p);
+    }
+
+    __syncthreads();
+
+    for(int i=threadIdx.x; i<task_num; i+=blockDim.x){
+        int g = group[i];
+        output[i] = priority[i]!=temp[g];
+    }
+};
+
+template<typename _Tg, typename _Tp>
+__global__ void cuda_do_task_group_priority(
+        const torch::PackedTensorAccessor<_Tg,2,torch::RestrictPtrTraits> group,
+        const torch::PackedTensorAccessor<_Tp,2,torch::RestrictPtrTraits> priority,
+        const torch::PackedTensorAccessor<bool,2,torch::RestrictPtrTraits> value,
+        torch::PackedTensorAccessor<bool,2,torch::RestrictPtrTraits> result,
+        const _Tg NG)
+{
+    const int NP = group.size(0);
+    const int NT = group.size(1);
+    const int p = blockIdx.x * blockDim.x + threadIdx.x;
+    if(p < NP)
+    {
+        extern __shared__ char _temp[];
+        auto temp = reinterpret_cast<_Tp*>(_temp);
+        temp += (threadIdx.x * NG);
+        for(_Tg g=0; g<NG; g++){
+            temp[g] = std::numeric_limits<_Tp>::max();
+        }
+
+        for(int t=0; t<NT; t++){
+            if(value[p][t]){
+                continue;
+            }
+            _Tg g = group[p][t];
+            _Tp _p = priority[p][t];
+            if(_p < temp[g]){
+                temp[g] = _p;
+            }
+        }
+
+        for(int t=0; t<NT; t++){
+            _Tg g = group[p][t];
+            if(priority[p][t]==temp[g]){
+                result[p][t] = false;
+            }
+        }
+    }
+};
+
+
+
+void task_group_priority_cuda(
+        int* group, int* priority, bool* value, bool* output,
+        const int batch_size, const int task_num, const int group_num, const int device)
+{
+    const int shared_mem = group_num * sizeof(int);
+
+    GRL_CHECK_CUDA(cudaSetDevice(device));
+
+    task_group_priority_kernel<<<batch_size, 256, shared_mem>>>(
+        group, priority, value, output, batch_size, task_num, group_num);
+
+    GRL_CHECK_CUDA(cudaGetLastError());
+};
+

csrc/task_group_priority.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "./common.h"
+
+/**
+ * tasks are divided into groups,
+ * tasks in a group are visited by it's priority.
+ * the min priority value of unvisited tasks in a group is computed,
+ * output is false, if the task's priority equal the computed min priority, otherwise output is true
+ *
+ * group: task's group, shape is (batch_size, task_num)
+ * priority: task's priority, shape is (batch_size, task_num)
+ * value: task is visited or not, shape is (batch_size, task_num)
+ *
+ * output: the result, shape is (batch_size, task_num)
+ */
+auto task_group_priority(
+        const torch::Tensor& group,
+        const torch::Tensor& priority,
+        const torch::Tensor& value) -> torch::Tensor;
+
+void task_group_priority_cpu(
+        int* group, int* priority, bool* value, bool* ouput,
+        int batch_size, int task_num, int group_num);
+
+void task_group_priority_cuda(
+        int* group, int* priority, bool* value, bool* ouput,
+        int batch_size, int task_num, int group_num, int device);
+
+

csrc/task_group_split.cpp

@@ -0,0 +1,69 @@
+#include "task_group_split.h"
+
+void task_group_split_cpu(
+        int* group, bool* value, bool* output,
+        const int batch_size, const int task_num, const int group_num)
+{
+    auto temp = torch::make_unique<bool[]>(group_num);
+    for(int b=0; b<batch_size; b++)
+    {
+        for(int i=0; i<group_num; i++){
+            temp[i] = false;
+        }
+        
+        for(int i=0; i<task_num; i++){
+            if(value[i]){
+                int g = group[i];
+                temp[g] = true;
+            }
+        }
+        
+        output[b] = false;
+        for(int i=0; i<task_num; i++){
+            int g = group[i];
+            if(temp[g] && !value[i]){
+                output[b] = true;
+                break;
+            }
+        }
+        
+        group += task_num;
+        value += task_num;
+    }
+};
+
+
+auto task_group_split(
+    const Tensor& group, const Tensor& value) -> Tensor
+{
+    auto device = group.device();
+    const int batch_size = group.size(0);
+    const int task_num = group.size(1);
+    const int group_num = group.max().item<int>() + 1;
+    const int _group_num = group.min().item<int>();
+
+    GRL_CHECK(group_num <= task_num && _group_num >= 0, "group value error");
+
+    GRL_CHECK_TENSOR(group, device, false, false, batch_size, task_num);
+    GRL_CHECK_TENSOR(value, device, false, false, batch_size, task_num);
+
+    auto output = torch::zeros({batch_size}, torch::dtype(torch::kBool).device(device));
+
+    switch(device.type())
+    {
+        case torch::kCPU:
+            task_group_split_cpu(group.data_ptr<int>(), value.data_ptr<bool>(),
+                                 output.data_ptr<bool>(), batch_size, task_num, group_num); 
+            break;
+#ifdef CUDA_FOUND
+        case torch::kCUDA:
+            task_group_split_cuda(group.data_ptr<int>(), value.data_ptr<bool>(),
+                                  output.data_ptr<bool>(), batch_size, task_num, group_num, device.index());
+            break;
+#endif
+        default:
+            GRL_ERROR("unsupported device: %s", device.str().c_str());
+    }
+    
+    return output;
+};

csrc/task_group_split.cu

@@ -0,0 +1,53 @@
+#include "task_group_split.h"
+
+__global__ void task_group_split_kernel(
+        int* group, bool* value, bool* output,
+        const int batch_size, const int task_num, const int group_num)
+{
+    group += blockIdx.x * task_num;
+    value += blockIdx.x * task_num;
+    extern __shared__ bool temp[];
+    
+    __shared__ bool split;
+    if(threadIdx.x == 0) split = false;
+
+    for(int i=threadIdx.x; i<group_num; i+=blockDim.x)
+    {
+        temp[i] = false;
+    }
+
+    __syncthreads();
+
+    for(int i=threadIdx.x; i<task_num; i+=blockDim.x)
+    {
+        int g = group[i];
+        if(value[i]) temp[g] = true; 
+    }
+    
+    __syncthreads();
+    
+    for(int i=threadIdx.x; i<task_num; i+=blockDim.x)
+    {
+        int g = group[i];
+        if(temp[g] && !value[i]) split = true;
+    }
+    
+    __syncthreads();
+
+    if(threadIdx.x == 0) output[blockIdx.x] = split;
+};
+
+void task_group_split_cuda(
+        int* group, bool* value, bool* output,
+        const int batch_size, const int task_num, const int group_num, const int device)
+{
+    const int shared_mem = group_num * sizeof(bool);
+
+    GRL_CHECK_CUDA(cudaSetDevice(device));
+
+    task_group_split_kernel<<<batch_size, 256, shared_mem>>>(
+        group, value, output, batch_size, task_num, group_num);
+
+    GRL_CHECK_CUDA(cudaGetLastError());
+};
+

csrc/task_group_split.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "./common.h"
+
+/**
+ * tasks are divided into groups, 
+ * if tasks in a group are all visited or all not visited,
+ * output is is false, otherwise output is true
+ *
+ * group: task's group, shape is (batch_size, task_num)
+ * value: task is visited or not, shape is (batch_size, task_num)
+ *
+ * output: the result, shape is (batch_size,)
+ */
+auto task_group_split(const Tensor& group, const Tensor& value) -> Tensor;
+
+void task_group_split_cpu(
+        int* group, bool* value, bool* output,
+        const int batch_size, const int task_num, const int group_num);
+
+void task_group_split_cuda(
+        int* group, bool* value, bool* output,
+        const int batch_size, const int task_num, const int group_num, const int device);
+

examples/batching/batching.py

@@ -0,0 +1,165 @@
+import json
+
+from greedrl import Problem, Solver
+from greedrl.feature import *
+from greedrl.variable import *
+
+features = [local_feature('task_area'),
+            local_feature('task_roadway'),
+            local_feature('task_area_group'),
+            sparse_local_feature('task_item_id', 'task_item_num'),
+            sparse_local_feature('task_item_owner_id', 'task_item_num'),
+            variable_feature('worker_task_item'),
+            variable_feature('worker_used_roadway'),
+            variable_feature('worker_used_area')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_item_id'),
+             feature_variable('task_item_num'),
+             feature_variable('task_item_owner_id'),
+             feature_variable('task_area'),
+             feature_variable('task_area_group'),
+             feature_variable('task_load'),
+             feature_variable('task_group'),
+             worker_variable('worker_load_limit'),
+             worker_variable('worker_area_limit'),
+             worker_variable('worker_area_group_limit'),
+             worker_task_item('worker_task_item', item_id='task_item_id', item_num='task_item_num'),
+             worker_task_item('worker_task_item_owner', item_id='task_item_owner_id', item_num='task_item_num'),
+             worker_used_resource('worker_used_load', task_require='task_load'),
+             worker_used_resource('worker_used_area', task_require='task_area'),
+             worker_used_resource('worker_used_roadway', task_require='task_roadway'),
+             worker_used_resource('worker_used_area_group', task_require='task_area_group')]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_worker_end(self):
+        return self.worker_used_load < self.worker_load_limit
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand_now <= 0
+        # mask |= task_group_priority(self.task_group, self.task_out_stock_time, mask)
+
+        NT = self.task_item_id.size(1)
+        worker_task_item = self.worker_task_item[:, None, :]
+        worker_task_item = worker_task_item.expand(-1, NT, -1)
+        task_item_in_worker = worker_task_item.gather(2, self.task_item_id.long())
+        task_item_in_worker = (task_item_in_worker > 0) & (self.task_item_num > 0)
+
+        worker_task_item_owner = self.worker_task_item_owner[:, None, :]
+        worker_task_item_owner = worker_task_item_owner.expand(-1, NT, -1)
+        task_item_owner_in_worker = worker_task_item_owner.gather(2, self.task_item_owner_id.long())
+        task_item_owner_in_worker = (task_item_owner_in_worker > 0) & (self.task_item_num > 0)
+
+        # 同一个sku，不同货主，不能在一个拣选单
+        mask |= torch.any(task_item_in_worker & ~task_item_owner_in_worker, 2)
+
+        worker_load_limit = self.worker_load_limit - self.worker_used_load
+        mask |= (self.task_load > worker_load_limit[:, None])
+
+        task_area = self.task_area + self.worker_used_area[:, None, :]
+        task_area_num = task_area.clamp(0, 1).sum(2, dtype=torch.int32)
+        mask |= (task_area_num > self.worker_area_limit[:, None])
+
+        tak_area_group = self.task_area_group + self.worker_used_area_group[:, None, :]
+        tak_area_group_num = tak_area_group.clamp(0, 1).sum(2, dtype=torch.int32)
+        mask |= (tak_area_group_num > self.worker_area_group_limit[:, None])
+
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        area_num = self.worker_used_area.clamp(0, 1).sum(1)
+        roadway_num = self.worker_used_roadway.clamp(0, 1).sum(1)
+        item_num = self.worker_task_item.clamp(0, 1).sum(1)
+        penalty = (self.worker_load_limit - self.worker_used_load) * 10
+        return area_num * 100 + roadway_num * 10 + item_num + penalty
+
+
+def make_problem_from_json(data):
+    if isinstance(data, str):
+        data = json.loads(data)
+    problem = Problem()
+    problem.id = data["id"]
+    if 'uuid' in data:
+        problem.uuid = data["uuid"]
+
+    problem.task_item_id = torch.tensor(data["task_item_id"], dtype=torch.int32)
+    problem.task_item_owner_id = torch.tensor(data["task_item_owner_id"], dtype=torch.int32)
+    problem.task_item_num = torch.tensor(data["task_item_num"], dtype=torch.int32)
+    problem.task_area = torch.tensor(data["task_area"], dtype=torch.int32)
+    problem.task_roadway = torch.tensor(data["task_roadway"], dtype=torch.int32)
+    problem.task_out_stock_time = torch.tensor(data["task_out_stock_time"], dtype=torch.int32)
+    problem.task_area_group = torch.tensor(data["task_area_group"], dtype=torch.int32)
+
+    NT = problem.task_item_id.size(0)
+    problem.task_load = torch.ones(NT, dtype=torch.int32)
+    problem.task_group = torch.zeros(NT, dtype=torch.int32)
+    problem.task_demand = torch.ones(NT, dtype=torch.int32)
+
+    problem.worker_load_limit = torch.tensor(data["worker_load_limit"], dtype=torch.int32)
+    problem.worker_area_limit = torch.tensor(data["worker_area_limit"], dtype=torch.int32)
+    problem.worker_area_group_limit = torch.tensor(data["worker_area_group_limit"], dtype=torch.int32)
+
+    problem.features = features
+    problem.variables = variables
+    problem.constraint = Constraint
+    problem.objective = Objective
+
+    return problem
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    assert batch_size == 1
+
+    NT = task_count
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem()
+        problem.id = i
+
+        device = Solver().device
+        p = torch.ones(NT, 1000, dtype=torch.float32, device=device)
+        problem.task_item_id = torch.multinomial(p, 10).to(torch.int32).cpu()
+        problem.task_item_owner_id = torch.multinomial(p, 10).to(torch.int32).cpu()
+        problem.task_item_num = torch.randint(0, 5, (NT, 10), dtype=torch.int32)
+        problem.task_area = torch.randint(0, 5, (NT, 10), dtype=torch.int32).clamp(0, 1)
+        problem.task_roadway = torch.randint(0, 5, (NT, 200), dtype=torch.int32).clamp(0, 1)
+        problem.task_area_group = torch.randint(0, 5, (NT, 10), dtype=torch.int32).clamp(0, 1)
+
+        problem.task_load = torch.ones(NT, dtype=torch.int32)
+        problem.task_group = torch.zeros(NT, dtype=torch.int32)
+        problem.task_demand = torch.ones(NT, dtype=torch.int32)
+
+        problem.worker_load_limit = torch.tensor([20], dtype=torch.int32)
+        problem.worker_area_limit = torch.tensor([10], dtype=torch.int32)
+        problem.worker_area_group_limit = torch.tensor([10], dtype=torch.int32)
+
+        problem.features = features
+        problem.variables = variables
+        problem.constraint = Constraint
+        problem.objective = Objective
+
+        problem_list.append(problem)
+
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

examples/cvrp/cvrp.py

@@ -0,0 +1,88 @@
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl import Problem
+
+features = [continuous_feature('task_demand'),
+            continuous_feature('worker_weight_limit'),
+            continuous_feature('distance_matrix'),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             worker_variable('worker_weight_limit'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand_now <= 0
+        # 车辆容量限制
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_demand_now * self.task_weight > worker_weight_limit[:, None]
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        return self.distance_last_to_this
+
+    def step_task(self):
+        return self.distance_last_to_this
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    assert task_count in (100, 1000, 2000, 5000)
+
+    weight_limit = 50
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem(True)
+        problem.id = torch.arange(batch_size) + i * batch_size;
+
+        problem.worker_weight_limit = torch.full((batch_size, 1), weight_limit, dtype=torch.int32)
+
+        N = task_count
+        problem.task_demand = torch.randint(1, 10, (batch_size, N), dtype=torch.int32)
+        problem.task_demand_x = problem.task_demand.float() / weight_limit
+
+        # 一个单位的task_demand的重量
+        problem.task_weight = torch.ones(batch_size, N, dtype=torch.int32)
+
+        loc = torch.rand(batch_size, N + 1, 2, dtype=torch.float32)
+        problem.task_location = loc[:, 1:, :]
+        problem.worker_location = loc[:, 0:1, :]
+
+        distance_matrix = torch.norm(loc[:, :, None, :] - loc[:, None, :, :], dim=3)
+        problem.distance_matrix = distance_matrix
+
+        problem.features = features
+        problem.variables = variables
+        problem.constraint = Constraint
+        problem.objective = Objective
+
+        problem_list.append(problem)
+
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

examples/cvrp/orts.py

@@ -0,0 +1,107 @@
+import sys
+import time
+import torch
+import argparse
+import utils
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor
+
+from ortools.constraint_solver import pywrapcp
+from ortools.constraint_solver import routing_enums_pb2
+
+
+def solve(problem, i, max_time):
+    scale = 100000
+    size = problem.task_demand.size(1)
+    demand = [0] + problem.task_demand[i].tolist()
+    capacity = problem.worker_weight_limit[i].tolist()
+    distance = (problem.distance_matrix[i] * scale + 0.5).to(torch.int32).tolist()
+
+    queue = mp.Queue()
+    p = mp.Process(target=do_solve, args=(size, demand, capacity, distance, max_time, queue))
+    p.start()
+    p.join()
+
+    return queue.get() / scale, queue.get()
+
+
+def do_solve(size, demand, capacity, distance, max_time, queue):
+    capacity = capacity * size
+
+    manager = pywrapcp.RoutingIndexManager(size + 1, size, 0)
+    routing = pywrapcp.RoutingModel(manager)
+
+    def distance_callback(from_index, to_index):
+        from_node = manager.IndexToNode(from_index)
+        to_node = manager.IndexToNode(to_index)
+        return distance[from_node][to_node]
+
+    distance_callback_index = routing.RegisterTransitCallback(distance_callback)
+    routing.SetArcCostEvaluatorOfAllVehicles(distance_callback_index)
+
+    def demand_callback(from_index):
+        from_node = manager.IndexToNode(from_index)
+        return demand[from_node]
+
+    demand_callback_index = routing.RegisterUnaryTransitCallback(demand_callback)
+    routing.AddDimensionWithVehicleCapacity(demand_callback_index, 0, capacity, True, 'capacity')
+
+    params = pywrapcp.DefaultRoutingSearchParameters()
+    params.first_solution_strategy = (routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC)
+    params.local_search_metaheuristic = (routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
+    params.time_limit.seconds = max_time
+
+    start_time = time.time()
+    solution = routing.SolveWithParameters(params)
+    spent_time = time.time() - start_time
+
+    queue.put(solution.ObjectiveValue())
+    queue.put(spent_time)
+
+
+def run_orts(task, max_time):
+    problem, i = task
+    return solve(problem, i, max_time)
+
+
+def main(args):
+    print("args: {}".format(vars(args)))
+    problem_size = args.problem_size
+    problem_count = args.problem_count
+    batch_size = args.batch_size
+
+    assert problem_count % batch_size == 0
+    batch_count = problem_count // batch_size
+    problem_list = utils.make_problem(batch_count, batch_size, problem_size)
+
+    executor = ThreadPoolExecutor(max_workers=args.threads)
+    task_list = [(p, i) for p in problem_list for i in range(batch_size)]
+
+    total_cost = 0
+    total_time = 0
+    for cost, elapse in executor.map(run_orts, task_list, [args.max_time] * problem_count):
+        total_cost += cost
+        total_time += elapse
+
+    avg_cost = total_cost / problem_count
+    avg_time = total_time / problem_count
+    print()
+    print("-----------------------------------------------------")
+    print("avg_cost: {:.4f}".format(avg_cost))
+    print("avg_time: {:.6f}s".format(avg_time))
+    print("total_count: {}".format(problem_count))
+    print("-----------------------------------------------------\n")
+    sys.stdout.flush()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--threads', default=20, type=int, help='number of threads')
+    parser.add_argument('--max_time', default=60, type=int, help='the time limit for the search in seconds')
+
+    parser.add_argument('--problem_size', default=100, type=int, choices=[100, 1000, 2000, 5000],  help='problem size')
+    parser.add_argument('--problem_count', default=128, type=int,  help='total number of generated problem instances')
+    parser.add_argument('--batch_size', default=128, type=int,  help='batch size for feedforwarding')
+
+    args = parser.parse_args()
+    main(args)

examples/cvrp/solve.py

@@ -0,0 +1,65 @@
+import os
+import sys
+import time
+import torch
+import argparse
+import utils
+from greedrl import Solver
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+
+def do_solve(args):
+    print("args: {}".format(vars(args)))
+
+    problem_size = args.problem_size
+    problem_count = args.problem_count
+    batch_size = args.batch_size
+    assert problem_count % batch_size == 0
+    batch_count = problem_count // batch_size
+
+    problem_list = utils.make_problem(batch_count, batch_size, problem_size)
+
+    solver = Solver(device=args.device)
+
+    model_path = os.path.join('./', args.model_name)
+    solver.load_agent(model_path)
+
+    total_cost = 0
+
+    if solver.device.type == 'cuda':
+        torch.cuda.synchronize()
+
+    start_time = time.time()
+    for problem in problem_list:
+        solution = solver.solve(problem, greedy=False, batch_size=batch_size)
+        total_cost += solution.cost.sum().item()
+
+    if solver.device.type == 'cuda':
+        torch.cuda.synchronize()
+
+    total_time = time.time() - start_time
+
+    avg_cost = total_cost / problem_count
+    avg_time = total_time / problem_count
+    print()
+    print("-----------------------------------------------------")
+    print("avg_cost: {:.4f}".format(avg_cost))
+    print("avg_time: {:.6f}s".format(avg_time))
+    print("total_count: {}".format(problem_count))
+    print("-----------------------------------------------------\n")
+    sys.stdout.flush()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', default='cpu', choices=['cpu', 'cuda'], help="choose a device")
+    parser.add_argument('--model_name', default='cvrp_100.pt', choices=['cvrp_100.pt', 'cvrp_1000.pt', 'cvrp_2000.pt', 'cvrp_5000.pt'], help="choose a model")
+    parser.add_argument('--problem_size', default=100, type=int, choices=[100, 1000, 2000, 5000],  help='problem size')
+    parser.add_argument('--problem_count', default=128, type=int,  help='total number of generated problem instances')
+    parser.add_argument('--batch_size', default=128, type=int,  help='batch size for feedforwarding')
+
+    args = parser.parse_args()
+    do_solve(args)
+

examples/cvrp/train.py

@@ -0,0 +1,83 @@
+import sys
+import math
+import argparse
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import utils
+from greedrl import Solver
+
+
+def do_train(args, rank):
+    world_size = args.world_size
+    model_filename = args.model_filename
+    problem_size = args.problem_size
+    batch_size = args.batch_size
+
+    index = model_filename.rfind('.')
+    if world_size > 1:
+        stdout_filename = '{}_r{}.log'.format(model_filename[0:index], rank)
+    else:
+        stdout_filename = '{}.log'.format(model_filename[0:index])
+
+    stdout = open(stdout_filename, 'a')
+    sys.stdout = stdout
+    sys.stderr = stdout
+
+    print("args: {}".format(vars(args)))
+    if world_size > 1:
+        dist.init_process_group('NCCL', init_method='tcp://127.0.0.1:29500',
+                                rank=rank, world_size=world_size)
+
+    problem_batch_size = 8
+    batch_count = 0
+    if problem_size == 100:
+        batch_count = math.ceil(10000 / problem_batch_size)
+    elif problem_size == 1000:
+        batch_count = math.ceil(200 / problem_batch_size)
+    elif problem_size == 2000:
+        batch_count = math.ceil(100 / problem_batch_size)
+    elif problem_size == 5000:
+        batch_count = math.ceil(10 / problem_batch_size)
+    else:
+        raise Exception("unsupported problem size: {}".format(problem_size))
+
+    nn_args = {
+        'encode_norm': 'instance',
+        'encode_layers': 6,
+        'decode_rnn': 'LSTM'
+    }
+
+    device = None if world_size == 1 else 'cuda:{}'.format(rank)
+    solver = Solver(device, nn_args)
+
+    train_dataset = utils.Dataset(None, problem_batch_size, problem_size)
+    valid_dataset = utils.Dataset(batch_count, problem_batch_size, problem_size)
+
+    solver.train(model_filename, train_dataset, valid_dataset,
+                 train_dataset_workers=5,
+                 batch_size=batch_size,
+                 memopt=10,
+                 topk_size=1,
+                 init_lr=1e-4,
+                 valid_steps=500,
+                 warmup_steps=0)
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')
+    parser.add_argument('--model_filename', type=str, help='model file name')
+    parser.add_argument('--problem_size', default=100, type=int, choices=[100, 1000, 2000, 5000],  help='problem size')
+    parser.add_argument('--batch_size', default=128, type=int,  help='batch size for training')
+
+    args = parser.parse_args()
+
+    processes = []
+    for rank in range(args.world_size):
+        p = mp.Process(target=do_train, args=(args, rank))
+        p.start()
+        processes.append(p)
+
+    for p in processes:
+        p.join()

examples/cvrp/utils.py

@@ -0,0 +1,65 @@
+from greedrl.feature import *
+from cvrp import make_problem as make_cvrp_problem
+from torch.utils.data import Dataset, IterableDataset, DataLoader
+
+
+def make_problem(batch_count, batch_size, task_count):
+    features = [continuous_feature('task_demand_x'),
+                continuous_feature('distance_matrix')]
+
+    problem_list = make_cvrp_problem(batch_count, batch_size, task_count)
+    for problem in problem_list:
+        problem.features = features
+
+    return problem_list
+
+
+class Dataset(IterableDataset):
+    def __init__(self, batch_count, batch_size, task_count):
+        self._batch_size = batch_size
+        self._task_count = task_count
+        self._batch_count = batch_count
+        self._index = 0
+
+    def __iter__(self):
+        self._index = 0
+        return self
+
+    def __next__(self):
+        if self._batch_count is not None \
+                and self._index >= self._batch_count:
+            raise StopIteration()
+
+        p = make_problem(1, self._batch_size, self._task_count)[0]
+        self._index += 1
+        return p
+
+
+def write_vrplib(filename, name, size, demand, capacity, location):
+    with open(filename, 'w') as f:
+        f.write('\n'.join([
+            "{} : {}".format(k, v)
+            for k, v in (
+                ('NAME', name),
+                ('TYPE', 'CVRP'),
+                ('COMMENT', 'NONE'),
+                ('DIMENSION', size + 1),
+                ('EDGE_WEIGHT_TYPE', 'EUC_2D'),
+                ('CAPACITY', capacity)
+            )
+        ]))
+
+        f.write('\n')
+        f.write('NODE_COORD_SECTION\n')
+
+        f.write('\n'.join(['{}\t{}\t{}'.format(i + 1, x, y) for i, (x, y) in enumerate(location)]))
+
+        f.write('\n')
+        f.write('DEMAND_SECTION\n')
+        f.write('\n'.join(['{}\t{}'.format(i + 1, d) for i, d in enumerate([0] + demand)]))
+
+        f.write('\n')
+        f.write('DEPOT_SECTION\n')
+        f.write('1\n')
+        f.write('-1\n')
+        f.write('EOF\n')

examples/dpdp/dpdp.py

@@ -0,0 +1,191 @@
+import json
+
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl.function import *
+from greedrl import Problem
+
+features = [local_category('task_order'),
+            global_category('task_type', 2),
+            global_category('task_new_order', 2),
+            variable_feature('time_this_to_task'),
+            continuous_feature('x_time_matrix'),
+            continuous_feature('task_due_time_x'),
+            continuous_feature('worker_task_mask')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             task_variable('task_pickup_this', feature='task_pickup'),
+             task_variable('task_due_time_this', feature='task_due_time'),
+             feature_variable('task_order', feature='task_order'),
+             feature_variable('task_type', feature='task_type'),
+             feature_variable('task_new_pickup', feature='task_new_pickup'),
+             feature_variable('worker_task_mask', feature='worker_task_mask'),
+             worker_count_now('worker_count_now', feature='worker_count'),
+             worker_variable('worker_min_old_task_this', feature='worker_min_old_task'),
+             worker_variable('worker_max_new_order_this', feature='worker_max_new_order'),
+             worker_variable('worker_task_mask_this', feature='worker_task_mask'),
+             worker_used_resource('worker_used_old_task', task_require='task_old'),
+             worker_used_resource('worker_used_new_order', task_require='task_new_pickup'),
+             worker_used_resource('worker_used_time', edge_require='time_matrix'),
+             edge_variable('time_this_to_task', feature='x_time_matrix', this_to_task=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_worker_start(self):
+        mask = self.worker_count_now <= 0
+
+        finished = self.task_demand_now <= 0
+        worker_task_mask = self.worker_task_mask | finished[:, None, :]
+        mask |= torch.all(worker_task_mask, 2)
+
+        return mask
+
+    def mask_worker_end(self):
+        mask = self.worker_used_old_task < self.worker_min_old_task_this
+        mask |= task_group_split(self.task_order, self.task_demand_now <= 0)
+        return mask
+
+    def mask_task(self):
+        mask = self.task_demand_now <= 0
+
+        mask |= task_group_priority(self.task_order, self.task_type, mask)
+
+        worker_max_new_order = self.worker_max_new_order_this - self.worker_used_new_order
+        mask |= self.task_new_pickup > worker_max_new_order[:, None]
+
+        mask |= self.worker_task_mask_this
+
+        return mask
+
+    def finished(self):
+        worker_mask = self.worker_count_now <= 0
+        task_mask = self.task_demand_now <= 0
+        worker_task_mask = worker_mask[:, :, None] | task_mask[:, None, :]
+
+        worker_task_mask |= self.worker_task_mask
+        batch_size = worker_task_mask.size(0)
+        worker_task_mask = worker_task_mask.view(batch_size, -1)
+        return worker_task_mask.all(1)
+
+
+class Objective:
+
+    def step_task(self):
+        over_time = (self.worker_used_time - self.task_due_time_this).clamp(min=0)
+        pickup_time = self.worker_used_time * self.task_pickup_this
+        return self.worker_used_time + over_time + pickup_time
+
+    def step_finish(self):
+        return self.task_demand_now.sum(1) * 1000
+
+
+def preprocess(problem):
+    NW, NT = problem.worker_task_mask.size()
+
+    worker_task_old = torch.ones(NW, NT, dtype=torch.int32)
+    new_task_mask = problem.task_new_order[None, :].expand(NW, NT)
+    worker_task_old[new_task_mask] = 0
+    worker_task_old[problem.worker_task_mask] = 0
+    assert torch.all(worker_task_old.sum(0) <= 1)
+    problem.worker_min_old_task = worker_task_old.sum(1)
+
+    problem.worker_count = torch.ones(NW, dtype=torch.int32)
+    problem.task_demand = torch.ones(NT, dtype=torch.int32)
+    problem.task_pickup = (problem.task_type == 0).to(torch.int32)
+
+    task_old = torch.ones(NT, dtype=torch.int32)
+    task_old[problem.task_new_order] = 0
+    problem.task_old = task_old
+
+    task_new_pickup = torch.ones(NT, dtype=torch.int32)
+    task_new_pickup[problem.task_type >= 1] = 0
+    task_new_pickup[~problem.task_new_order] = 0
+    problem.task_new_pickup = task_new_pickup
+
+    problem.task_due_time_x = problem.task_due_time.float() / 900
+    problem.x_time_matrix = problem.time_matrix.float() / 900
+
+    problem.features = features
+    problem.variables = variables
+    problem.constraint = Constraint
+    problem.objective = Objective
+
+    return problem
+
+
+def make_problem_from_json(data):
+    data = json.loads(data)
+
+    problem = Problem()
+
+    problem.id = data['id']
+    problem.task_order = torch.tensor(data['task_order'], dtype=torch.int32)
+    problem.task_type = torch.tensor(data['task_type'], dtype=torch.int32)
+    problem.task_new_order = torch.tensor(data['task_new_order'], dtype=torch.bool)
+    problem.task_due_time = torch.tensor(data['task_due_time'], dtype=torch.int32)
+
+    problem.worker_max_new_order = torch.tensor(data['worker_max_new_order'], dtype=torch.int32)
+    problem.worker_task_mask = torch.tensor(data['worker_task_mask'], dtype=torch.bool)
+    problem.time_matrix = torch.tensor(data['time_matrix'], dtype=torch.int32)
+
+    NW, NT = problem.worker_task_mask.size()
+
+    assert problem.task_order.size() == (NT,), "task_order size error"
+    assert problem.task_type.size() == (NT,), "task_type size error"
+    assert problem.task_new_order.size() == (NT,), "task_new_order size error"
+    assert problem.task_due_time.size() == (NT,), "task_due_time size error"
+    assert problem.worker_max_new_order.size() == (NW,), "worker_max_new_order size error"
+    assert problem.time_matrix.size() == (NW + NT, NW + NT), "time_matrix size error"
+
+    return preprocess(problem)
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    assert batch_size == 1
+    assert task_count == 100
+
+    NW = 100
+    NT = task_count
+    NO = NT // 2  # 订单数, 一个订单有pickup， delivery两个任务
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem()
+
+        # user-provided data
+        problem.worker_max_new_order = torch.full((NW,), 2, dtype=torch.int32)
+
+        task_order = torch.arange(NO, dtype=torch.int32)
+        problem.task_order = torch.cat([task_order, task_order], 0)
+
+        task_type = torch.zeros(NO, dtype=torch.int32)
+        problem.task_type = torch.cat([task_type, task_type + 1], 0)
+
+        problem.task_new_order = torch.ones(NT, dtype=torch.bool)
+
+        task_due_time = torch.randint(1000, 1800, (NO,), dtype=torch.int32)
+        problem.task_due_time = torch.cat([task_due_time, task_due_time + 1800], 0)
+
+        worker_task_mask = torch.rand(NW, NO) < 0.9
+        problem.worker_task_mask = torch.cat([worker_task_mask, worker_task_mask], 1)
+
+        loc = torch.rand(NW + NT, 2, dtype=torch.float32)
+        time_matrix = torch.norm(loc[:, None, :] - loc[None, :, :], dim=2) * 1000
+        problem.time_matrix = time_matrix.to(torch.int32)
+
+        problem_list.append(preprocess(problem))
+
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

examples/pdptw/pdptw.py

@@ -0,0 +1,136 @@
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl.function import *
+from greedrl import Problem
+
+features = [local_category('task_group'),
+            global_category('task_priority', 2),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             feature_variable('task_group'),
+             feature_variable('task_priority'),
+             feature_variable('task_due_time2', feature='task_due_time'),
+             task_variable('task_due_time'),
+             task_variable('task_service_time'),
+             task_variable('task_due_time_penalty'),
+             worker_variable('worker_basic_cost'),
+             worker_variable('worker_distance_cost'),
+             worker_variable('worker_due_time'),
+             worker_variable('worker_weight_limit'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             worker_used_resource('worker_used_time', 'distance_matrix', 'task_service_time', 'task_ready_time',
+                                  'worker_ready_time'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_worker_end(self):
+        return task_group_split(self.task_group, self.task_demand_now <= 0)
+
+    def mask_task(self):
+        mask = self.task_demand_now <= 0
+        mask |= task_group_priority(self.task_group, self.task_priority, mask)
+
+        worker_used_time = self.worker_used_time[:, None] + self.distance_this_to_task
+        mask |= (worker_used_time > self.task_due_time2) & (self.task_priority == 0)
+
+        # 容量约束
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_demand_now * self.task_weight > worker_weight_limit[:, None]
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_start(self):
+        return self.worker_basic_cost
+
+    def step_worker_end(self):
+        feasible = self.worker_used_time <= self.worker_due_time
+        return self.distance_last_to_this * self.worker_distance_cost, feasible
+
+    def step_task(self):
+        worker_used_time = self.worker_used_time - self.task_service_time
+        feasible = worker_used_time <= self.task_due_time
+        feasible &= worker_used_time <= self.worker_due_time
+        cost = self.distance_last_to_this * self.worker_distance_cost
+        return torch.where(feasible, cost, cost + self.task_due_time_penalty), feasible
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    assert batch_size == 1
+
+    N = task_count // 2  # 订单数, 一个订单有pickup， delivery两个任务
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem()
+        problem.id = i
+
+        problem.worker_weight_limit = torch.tensor([50], dtype=torch.float32)
+        problem.worker_ready_time = torch.tensor([0], dtype=torch.float32)
+        problem.worker_due_time = torch.tensor([1000000], dtype=torch.float32)
+        problem.worker_basic_cost = torch.tensor([100], dtype=torch.float32)
+        problem.worker_distance_cost = torch.tensor([1], dtype=torch.float32)
+
+        task_demand = torch.randint(1, 10, (N,), dtype=torch.int32)
+        problem.task_demand = torch.cat([task_demand, task_demand], 0)
+
+        task_weight = torch.ones(N, dtype=torch.float32)
+        problem.task_weight = torch.cat([task_weight, task_weight * -1], 0)
+
+        task_group = torch.arange(N, dtype=torch.int32)
+        problem.task_group = torch.cat([task_group, task_group], 0)
+
+        task_priority = torch.zeros(N, dtype=torch.int32)
+        problem.task_priority = torch.cat([task_priority, task_priority + 1], 0)
+
+        task_ready_time = torch.zeros(N, dtype=torch.float32)
+        problem.task_ready_time = torch.cat([task_ready_time, task_ready_time], 0)
+
+        task_due_time = torch.randint(10000, 100000, (N,), dtype=torch.float32)
+        problem.task_due_time = torch.cat([task_due_time, task_due_time * 2], 0)
+
+        task_service_time = torch.zeros(N, dtype=torch.float32)
+        problem.task_service_time = torch.cat([task_service_time, task_service_time])
+
+        task_due_time_penalty = torch.ones(N, dtype=torch.float32)
+        problem.task_due_time_penalty = torch.cat([task_due_time_penalty, task_due_time_penalty])
+
+        loc = torch.rand(N + 1, 2, dtype=torch.float32)
+        distance_matrix = torch.norm(loc[:, None, :] - loc[None, :, :], dim=2) * 1000
+        distance_matrix = distance_matrix.to(torch.float32)
+        index = torch.cat([torch.zeros(N + 1, dtype=torch.int64), torch.arange(N, dtype=torch.int64) + 1])
+        index1 = index[:, None]
+        index2 = index[None, :]
+        problem.distance_matrix = distance_matrix[index1, index2]
+
+        problem.features = features
+        problem.variables = variables
+        problem.constraint = Constraint
+        problem.objective = Objective
+
+        problem_list.append(problem)
+
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

examples/runner.py

@@ -0,0 +1,38 @@
+import time
+import random
+import argparse
+import torch
+
+from greedrl import Problem, Solution, Solver
+
+
+def run(make_problem, mask_task_ratio=0.1):
+    random.seed(123)
+    torch.manual_seed(123)
+    problem_list = make_problem(1)
+
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument('--device', default=None, type=str)
+    parser.add_argument('--batch_size', default=32, type=int)
+    parser.add_argument('--agent_file', default=None, type=str)
+    parser.add_argument('--valid_steps', default=5, type=int)
+    parser.add_argument('--max_steps', default=10000000, type=int)
+
+    args, _ = parser.parse_known_args()
+    for k, v in args.__dict__.items():
+        print("arg: {} = {}".format(k, v))
+
+    # rl train
+    solver = Solver(device=args.device)
+    solver.train(args.agent_file, problem_list, problem_list,
+                 batch_size=args.batch_size, valid_steps=args.valid_steps, max_steps=args.max_steps)
+    # predict
+    solver = Solver(device=args.device)
+    if args.agent_file is not None:
+        solver.load_agent(args.agent_file)
+
+    print("solve ...")
+    start = time.time()
+    for problem in problem_list:
+        solver.solve(problem, batch_size=args.batch_size)
+    print("time: {}s".format(time.time() - start))

examples/sdvrp/sdvrp.py

@@ -0,0 +1,83 @@
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl import Problem
+
+features = [continuous_feature('task_demand'),
+            continuous_feature('worker_weight_limit'),
+            continuous_feature('distance_matrix'),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             task_variable('task_weight_this', feature='task_weight'),
+             worker_variable('worker_weight_limit'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        return torch.min(self.task_demand_this, worker_weight_limit // self.task_weight_this)
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand <= 0
+        # 车辆容量限制
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        # 至少要能装下一个单位的demand
+        mask |= self.task_weight > worker_weight_limit[:, None]
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        return self.distance_last_to_this
+
+    def step_task(self):
+        return self.distance_last_to_this
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    assert batch_size == 1
+
+    NT = task_count
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem()
+        problem.id = i
+
+        problem.worker_weight_limit = [50]
+
+        problem.task_demand = torch.randint(1, 10, (NT,), dtype=torch.int64)
+
+        # 一个单位的task_demand的重量
+        problem.task_weight = torch.ones(NT, dtype=torch.int64)
+
+        loc = torch.rand(NT + 1, 2, dtype=torch.float32)
+        distance_matrix = torch.norm(loc[:, None, :] - loc[None, :, :], dim=2) * 1000
+        problem.distance_matrix = distance_matrix.to(torch.int64)
+
+        problem.variables = variables
+        problem.constraint = Constraint
+        problem.objective = Objective
+
+        problem_list.append(problem)
+
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

examples/tsp/tsp.py

@@ -0,0 +1,74 @@
+from greedrl.feature import *
+from greedrl.variable import *
+from greedrl import Problem
+
+features = [continuous_feature('task_location'),
+            variable_feature('distance_this_to_task'),
+            variable_feature('distance_task_to_end')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True),
+             edge_variable('distance_last_to_loop', feature='distance_matrix', last_to_loop=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand_now <= 0
+        return mask
+
+    def mask_worker_end(self):
+        return torch.any(self.task_demand_now > 0, 1)
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_end(self):
+        return self.distance_last_to_loop
+
+    def step_task(self):
+        return self.distance_last_to_this
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    NP = batch_size
+    NT = task_count
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem(True)
+
+        problem.task_demand = torch.ones(NP, NT, dtype=torch.int32)
+
+        loc = torch.rand(NP, NT + 1, 2, dtype=torch.float32)
+        problem.distance_matrix = torch.norm(loc[:, :, None, :] - loc[:, None, :, :], dim=3)
+        problem.distance_matrix[0, :] = 0
+        problem.distance_matrix[:, 0] = 0
+
+        problem.task_location = loc[:, 1:]
+
+        problem.features = features
+        problem.variables = variables
+        problem.constraint = Constraint
+        problem.objective = Objective
+
+        problem_list.append(problem)
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

examples/vrptw/vrptw.py

@@ -0,0 +1,141 @@
+import json
+
+from greedrl import Problem
+from greedrl.feature import *
+from greedrl.variable import *
+
+features = [continuous_feature('worker_weight_limit'),
+            continuous_feature('worker_ready_time'),
+            continuous_feature('worker_due_time'),
+            continuous_feature('worker_basic_cost'),
+            continuous_feature('worker_distance_cost'),
+            continuous_feature('task_demand'),
+            continuous_feature('task_weight'),
+            continuous_feature('task_ready_time'),
+            continuous_feature('task_due_time'),
+            continuous_feature('task_service_time'),
+            continuous_feature('distance_matrix')]
+
+variables = [task_demand_now('task_demand_now', feature='task_demand'),
+             task_demand_now('task_demand_this', feature='task_demand', only_this=True),
+             feature_variable('task_weight'),
+             feature_variable('task_due_time'),
+             feature_variable('task_ready_time'),
+             feature_variable('task_service_time'),
+             worker_variable('worker_weight_limit'),
+             worker_variable('worker_due_time'),
+             worker_variable('worker_basic_cost'),
+             worker_variable('worker_distance_cost'),
+             worker_used_resource('worker_used_weight', task_require='task_weight'),
+             worker_used_resource('worker_used_time', 'distance_matrix', 'task_service_time', 'task_ready_time',
+                                  'worker_ready_time'),
+             edge_variable('distance_last_to_this', feature='distance_matrix', last_to_this=True),
+             edge_variable('distance_this_to_task', feature='distance_matrix', this_to_task=True),
+             edge_variable('distance_task_to_end', feature='distance_matrix', task_to_end=True)]
+
+
+class Constraint:
+
+    def do_task(self):
+        return self.task_demand_this
+
+    def mask_task(self):
+        # 已经完成的任务
+        mask = self.task_demand_now <= 0
+        # 车辆容量限制
+        worker_weight_limit = self.worker_weight_limit - self.worker_used_weight
+        mask |= self.task_demand_now * self.task_weight > worker_weight_limit[:, None]
+
+        worker_used_time = self.worker_used_time[:, None] + self.distance_this_to_task
+        mask |= worker_used_time > self.task_due_time
+
+        worker_used_time = torch.max(worker_used_time, self.task_ready_time)
+        worker_used_time += self.task_service_time
+        worker_used_time += self.distance_task_to_end
+        mask |= worker_used_time > self.worker_due_time[:, None]
+
+        return mask
+
+    def finished(self):
+        return torch.all(self.task_demand_now <= 0, 1)
+
+
+class Objective:
+
+    def step_worker_start(self):
+        return self.worker_basic_cost
+
+    def step_worker_end(self):
+        return self.distance_last_to_this * self.worker_distance_cost
+
+    def step_task(self):
+        return self.distance_last_to_this * self.worker_distance_cost
+
+
+def make_problem_from_json(data):
+    if isinstance(data, str):
+        data = json.loads(data)
+
+    problem = Problem()
+    problem.worker_weight_limit = torch.tensor(data['worker_weight_limit'], dtype=torch.float32)
+    problem.worker_ready_time = torch.tensor(data['worker_ready_time'], dtype=torch.float32)
+    problem.worker_due_time = torch.tensor(data['worker_due_time'], dtype=torch.float32)
+    problem.worker_basic_cost = torch.tensor(data['worker_basic_cost'], dtype=torch.float32)
+    problem.worker_distance_cost = torch.tensor(data['worker_distance_cost'], dtype=torch.float32)
+
+    problem.task_demand = torch.tensor(data['task_demand'], dtype=torch.int32)
+    problem.task_weight = torch.tensor(data['task_weight'], dtype=torch.float32)
+    problem.task_ready_time = torch.tensor(data['task_ready_time'], dtype=torch.float32)
+    problem.task_due_time = torch.tensor(data['task_due_time'], dtype=torch.float32)
+    problem.task_service_time = torch.tensor(data['task_service_time'], dtype=torch.float32)
+
+    problem.distance_matrix = torch.tensor(data['distance_matrix'], dtype=torch.float32);
+
+    problem.features = features
+    problem.variables = variables
+    problem.constraint = Constraint
+    problem.objective = Objective
+
+    return problem
+
+
+def make_problem(batch_count, batch_size=1, task_count=100):
+    assert batch_size == 1
+
+    NT = task_count
+    problem_list = []
+    for i in range(batch_count):
+        problem = Problem()
+        problem.id = i
+
+        problem.worker_weight_limit = torch.tensor([50], dtype=torch.float32)
+        problem.worker_ready_time = torch.tensor([0], dtype=torch.float32)
+        problem.worker_due_time = torch.tensor([1000000], dtype=torch.float32)
+        problem.worker_basic_cost = torch.tensor([100], dtype=torch.float32)
+        problem.worker_distance_cost = torch.tensor([1], dtype=torch.float32)
+
+        problem.task_demand = torch.randint(1, 10, (NT,), dtype=torch.int32)
+        problem.task_weight = torch.ones(NT, dtype=torch.float32)
+        problem.task_ready_time = torch.zeros(NT, dtype=torch.float32)
+        problem.task_due_time = torch.randint(10000, 100000, (NT,), dtype=torch.float32)
+        problem.task_service_time = torch.zeros(NT, dtype=torch.float32)
+
+        loc = torch.rand(NT + 1, 2, dtype=torch.float32)
+        problem.distance_matrix = torch.norm(loc[:, None, :] - loc[None, :, :], dim=2) * 1000
+        problem_list.append(problem)
+
+        problem.features = features
+        problem.variables = variables
+        problem.constraint = Constraint
+        problem.objective = Objective
+
+    return problem_list
+
+
+if __name__ == '__main__':
+    import sys
+    import os.path as osp
+    sys.path.append(osp.join(osp.dirname(__file__), '../'))
+    import runner
+
+    runner.run(make_problem)

greedrl/.gitignore

@@ -0,0 +1,2 @@
+*.c
+version.py

greedrl/__init__.py

@@ -0,0 +1,8 @@
+import sys
+
+from .solver import Problem, Solution, Solver
+from .const import GRL_WORKER_START, GRL_WORKER_END, GRL_TASK, GRL_FINISH
+
+
+greedrl = sys.modules[__name__]
+

greedrl/agent.py

@@ -0,0 +1,203 @@
+import torch
+
+from torch import nn
+from collections import OrderedDict
+from torch.utils.checkpoint import checkpoint
+from .feature import *
+from .pyenv import PyEnv
+from .encode import Encode
+from .decode import Decode
+
+
+class Agent(nn.Module):
+
+    def __init__(self, nn_args):
+        super(Agent, self).__init__()
+
+        self.nn_args = nn_args
+        self.vars_dim = sum(nn_args['variable_dim'].values())
+        self.steps_ratio = nn_args.setdefault('decode_steps_ratio', 1.0);
+
+        logit_clips = nn_args.setdefault('decode_logit_clips', 10.0);
+        if isinstance(logit_clips, str):
+            self.logit_clips = [float(v) for v in logit_clips.split(',')]
+        else:
+            self.logit_clips = [float(logit_clips)]
+
+        self.nn_encode = Encode(nn_args)
+        self.nn_decode = Decode(nn_args)
+
+    def nn_args_dict(self):
+        return self.nn_args
+
+    def forward(self, problem, batch_size, greedy=False, solution=None, memopt=0):
+        X, K, V = self.nn_encode(problem.feats, problem.batch_size,
+                                 problem.worker_num, problem.task_num, memopt)
+
+        return self.interact(problem, X, K, V, batch_size, greedy, solution, memopt)
+
+    def interact(self, problem, X, K, V, batch_size, greedy, solution, memopt):
+        NP = problem.batch_size
+        NW = problem.worker_num
+        NT = problem.task_num
+
+        sample_num = batch_size // NP
+        assert sample_num > 0 and batch_size % NP == 0
+
+        MyEnv = problem.environment
+        if MyEnv is None:
+            env = PyEnv(problem, batch_size, sample_num, self.nn_args)
+        else:
+            env = MyEnv(str(problem.device), problem.feats, batch_size,
+                        sample_num, problem.worker_num, problem.task_num)
+
+        query = X.new_zeros(batch_size, X.size(-1))
+        state1 = X.new_zeros(batch_size, X.size(-1))
+        state2 = X.new_zeros(batch_size, X.size(-1))
+
+        p_list = []
+        NULL = X.new_ones(0)
+        p_index = torch.div(torch.arange(batch_size, device=X.device), sample_num, rounding_mode='trunc') # torch.arange(batch_size, device=X.device) // sample_num
+        if solution is not None:
+            solution = solution[:, :, 0:2].to(torch.int64).permute(1, 0, 2)
+            assert torch.all(solution >= 0) and solution.size(1) == batch_size
+            offset = torch.tensor([0, NW, NW + NW, NW + NW + NT], device=X.device)
+            chosen_list = solution[:, :, 1] + offset[solution[:, :, 0]]
+
+            mode = 0
+            sample_p = torch.rand(batch_size, device=X.device)
+            for chosen in chosen_list:
+                env_time = env.time()
+                clip = self.logit_clips[min(env_time, len(self.logit_clips) - 1)]
+                varfeat = env.make_feat() if self.vars_dim > 0 else NULL
+                state1, state2, chosen_p = self.decode(X, K, V, query, state1, state2,
+                                                       varfeat, env.mask(), chosen, sample_p, clip, mode, memopt)
+                query = X[p_index, chosen]
+                p_list.append(chosen_p)
+                env.step(chosen)
+
+            assert env.all_finished(), 'not all finished!'
+        else:
+            mode = 1 if greedy else 2
+            min_env_time = int(self.steps_ratio * NT)
+            R = torch.rand(NT * 2, batch_size, device=X.device)
+            while True:
+                env_time = env.time()
+                if env_time > min_env_time and env_time % 3 == 0 and env.all_finished():
+                    break
+
+                clip = self.logit_clips[min(env_time, len(self.logit_clips) - 1)]
+                sample_p = R[env_time % R.size(0)]
+                chosen = X.new_empty(batch_size, dtype=torch.int64)
+                varfeat = env.make_feat() if self.vars_dim > 0 else NULL
+                state1, state2, chosen_p = self.decode(X, K, V, query, state1, state2,
+                                                       varfeat, env.mask(), chosen, sample_p, clip, mode, memopt)
+                query = X[p_index, chosen]
+                p_list.append(chosen_p)
+                env.step(chosen)
+
+        env.finalize()
+        return env, torch.stack(p_list, 1)
+
+    def decode(self, X, K, V, query, state1, state2, varfeat, mask, chosen, sample_p, clip, mode, memopt):
+        run_fn = self.decode_fn(clip, mode, memopt)
+        if self.training and memopt > 3:
+            return checkpoint(run_fn, X, K, V, query, state1, state2, varfeat, mask, chosen, sample_p)
+        else:
+            return run_fn(X, K, V, query, state1, state2, varfeat, mask, chosen, sample_p)
+
+    def decode_fn(self, clip, mode, memopt):
+        memopt = 0 if memopt > 3 else memopt
+
+        def run_fn(X, K, V, query, state1, state2, varfeat, mask, chosen, sample_p):
+            return self.nn_decode(X, K, V, query, state1, state2,
+                                  varfeat, mask, chosen, sample_p, clip, mode, memopt)
+
+        return run_fn
+
+
+def parse_nn_args(problem, nn_args):
+    worker_dim = OrderedDict()
+    task_dim = OrderedDict()
+    edge_dim = OrderedDict()
+    variable_dim = OrderedDict()
+    embed_dict = OrderedDict()
+
+    def set_dim_by_name(name, k, dim):
+        if name.startswith("worker_task_"):
+            edge_dim[k] = dim
+        elif name.startswith("worker_"):
+            worker_dim[k] = dim
+        elif name.startswith("task_"):
+            task_dim[k] = dim
+        elif name.endswith("_matrix"):
+            edge_dim[k] = dim
+        else:
+            raise Exception("attribute can't be feature: {}".format(k))
+
+    feature_dict = make_feat_dict(problem)
+    variables = [var(problem, problem.batch_size, 1) for var in problem.variables]
+    variable_dict = dict([(var.name, var) for var in variables])
+    for k, f in feature_dict.items():
+        if isinstance(f, VariableFeature):
+            var = variable_dict[f.name]
+            assert hasattr(var, 'make_feat'), \
+                "{} cann't be variable feature, name:{}".format(type(var).__name__, k)
+            v = var.make_feat()
+            if v.dim() == 2:
+                variable_dim[k] = 1
+            else:
+                variable_dim[k] = v.size(-1)
+        elif isinstance(f, SparseLocalFeature):
+            edge_dim[k] = 1
+            set_dim_by_name(f.value, k, 1)
+        elif isinstance(f, LocalFeature):
+            edge_dim[k] = 1
+            set_dim_by_name(f.name, k, 1)
+        elif isinstance(f, LocalCategory):
+            edge_dim[k] = 1
+        elif isinstance(f, GlobalCategory):
+            set_dim_by_name(f.name, k, nn_args.setdefault('encode_hidden_dim', 128))
+            embed_dict[k] = f.size
+        elif isinstance(f, ContinuousFeature):
+            v = problem.feats[k]
+            if k.startswith("worker_task_") or k.endswith("_matrix"):
+                simple_dim = 3
+            else:
+                simple_dim = 2
+
+            if v.dim() == simple_dim:
+                set_dim_by_name(f.name, k, 1)
+            else:
+                set_dim_by_name(f.name, k, v.size(-1))
+        else:
+            raise Exception("unsupported feature type: {}".format(type(f)))
+
+    nn_args['worker_dim'] = worker_dim
+    nn_args['task_dim'] = task_dim
+    nn_args['edge_dim'] = edge_dim
+    nn_args['variable_dim'] = variable_dim
+    nn_args['embed_dict'] = embed_dict
+    nn_args['feature_dict'] = feature_dict
+    return nn_args
+
+
+def make_feat_dict(problem):
+    feature_dict = OrderedDict()
+
+    def add(k, f):
+        _f = feature_dict.get(k)
+        if _f is None or _f == f:
+            feature_dict[k] = f
+        else:
+            "duplicated feature, name: {}, feature1: {}, feature2: {}".format(k, _f, f)
+
+    for f in problem.features:
+        if isinstance(f, VariableFeature):
+            add(':'.join(['var', f.name]), f)
+        elif isinstance(f, SparseLocalFeature):
+            add(':'.join([f.index, f.value]), f)
+        else:
+            add(f.name, f)
+
+    return feature_dict

greedrl/const.py

@@ -0,0 +1,7 @@
+
+GRL_WORKER_START = 0
+GRL_WORKER_END = 1
+GRL_TASK = 2
+GRL_FINISH = 3
+
+

greedrl/decode.py

@@ -0,0 +1,196 @@
+import math
+import torch
+import torch.nn.functional as F
+
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, heads, hidden_dim):
+        super(MultiHeadAttention, self).__init__()
+
+        assert hidden_dim % heads == 0
+
+        self.heads = heads
+        head_dim = hidden_dim // heads
+        self.alpha = 1 / math.sqrt(head_dim)
+
+        self.nn_Q = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+        self.nn_O = nn.Parameter(torch.Tensor(hidden_dim, hidden_dim))
+
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, q, K, V, mask):
+        batch_size, query_num, hidden_dim = q.size()
+
+        size = (self.heads, batch_size, query_num, -1)
+
+        q = q.reshape(-1, hidden_dim)
+        Q = torch.matmul(q, self.nn_Q).view(size)
+
+        value_num = V.size(2)
+        heads_batch = self.heads * batch_size
+        Q = Q.view(heads_batch, query_num, -1)
+        K = K.view(heads_batch, value_num, -1).transpose(1, 2)
+
+        S = masked_tensor(mask, self.heads)
+        S = S.view(heads_batch, query_num, value_num)
+        S.baddbmm_(Q, K, alpha=self.alpha)
+        S = S.view(self.heads, batch_size, query_num, value_num)
+
+        S = F.softmax(S, dim=-1)
+
+        x = torch.matmul(S, V).permute(1, 2, 0, 3)
+        x = x.reshape(batch_size, query_num, -1)
+        x = torch.matmul(x, self.nn_O)
+        return x
+
+
+class Decode(nn.Module):
+
+    def __init__(self, nn_args):
+        super(Decode, self).__init__()
+
+        self.nn_args = nn_args
+
+        heads = nn_args['decode_atten_heads']
+        hidden_dim = nn_args['decode_hidden_dim']
+
+        self.heads = heads
+        self.alpha = 1 / math.sqrt(hidden_dim)
+
+        if heads > 0:
+            assert hidden_dim % heads == 0
+            head_dim = hidden_dim // heads
+            self.nn_K = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+            self.nn_V = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+            self.nn_mha = MultiHeadAttention(heads, hidden_dim)
+
+        decode_rnn = nn_args.setdefault('decode_rnn', 'LSTM')
+        assert decode_rnn in ('GRU', 'LSTM', 'NONE')
+        if decode_rnn == 'GRU':
+            self.nn_rnn_cell = nn.GRUCell(hidden_dim, hidden_dim)
+        elif decode_rnn == 'LSTM':
+            self.nn_rnn_cell = nn.LSTMCell(hidden_dim, hidden_dim)
+        else:
+            self.nn_rnn_cell = None
+
+        self.vars_dim = sum(nn_args['variable_dim'].values())
+        if self.vars_dim > 0:
+            atten_type = nn_args.setdefault('decode_atten_type', 'add')
+            assert atten_type == 'add', "must be addition attention when vars_dim > 0, {}".format(atten_type)
+            self.nn_A = nn.Parameter(torch.Tensor(self.vars_dim, hidden_dim))
+            self.nn_B = nn.Parameter(torch.Tensor(hidden_dim))
+        else:
+            atten_type = nn_args.setdefault('decode_atten_type', 'prod')
+
+        if atten_type == 'add':
+            self.nn_W = nn.Parameter(torch.Tensor(hidden_dim))
+        else:
+            self.nn_W = None
+
+        for param in self.parameters():
+            stdv = 1 / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, X, K, V, query, state1, state2, varfeat, mask, chosen, sample_p, clip, mode, memopt=0):
+        if self.training and memopt > 2:
+            state1, state2 = checkpoint(self.rnn_step, query, state1, state2)
+        else:
+            state1, state2 = self.rnn_step(query, state1, state2)
+
+        query = state1
+        NP = X.size(0)
+        NR = query.size(0) // NP
+        batch_size = query.size(0)
+        if self.heads > 0:
+            query = query.view(NP, NR, -1)
+            if self.training and memopt > 1:
+                query = checkpoint(self.nn_mha, query, K, V, mask)
+            else:
+                query = self.nn_mha(query, K, V, mask)
+
+            query = query.view(batch_size, -1)
+
+        if self.nn_W is None:
+            query = query.view(NP, NR, -1)
+            logit = masked_tensor(mask, 1)
+            logit = logit.view(NP, NR, -1)
+            X = X.permute(0, 2, 1)
+            logit.baddbmm_(query, X, alpha=self.alpha)
+            logit = logit.view(batch_size, -1)
+        else:
+            if self.training and self.vars_dim > 0 and memopt > 0:
+                logit = checkpoint(self.atten, query, X, varfeat, mask)
+            else:
+                logit = self.atten(query, X, varfeat, mask)
+
+        chosen_p = choose(logit, chosen, sample_p, clip, mode)
+        return state1, state2, chosen_p
+
+    def rnn_step(self, query, state1, state2):
+        if isinstance(self.nn_rnn_cell, nn.GRUCell):
+            state1 = self.nn_rnn_cell(query, state1)
+        elif isinstance(self.nn_rnn_cell, nn.LSTMCell):
+            state1, state2 = self.nn_rnn_cell(query, (state1, state2))
+        return state1, state2
+
+    def atten(self, query, keyvalue, varfeat, mask):
+        if self.vars_dim > 0:
+            varfeat = vfaddmm(varfeat, mask, self.nn_A, self.nn_B)
+        return atten(query, keyvalue, varfeat, mask, self.nn_W)
+
+
+def choose(logit, chosen, sample_p, clip, mode):
+    mask = logit == -math.inf
+    logit = torch.tanh(logit) * clip
+    logit[mask] = -math.inf
+
+    if mode == 0:
+        pass
+    elif mode == 1:
+        chosen[:] = logit.argmax(1)
+    elif mode == 2:
+        p = logit.exp()
+        chosen[:] = torch.multinomial(p, 1).squeeze(1)
+    else:
+        raise Exception()
+
+    logp = logit.log_softmax(1)
+    logp = logp.gather(1, chosen[:, None])
+    logp = logp.squeeze(1)
+    return logp
+
+
+def atten(query, keyvalue, varfeat, mask, weight):
+    batch_size = query.size(0)
+    NP, NK, ND = keyvalue.size()
+
+    query = query.view(NP, -1, 1, ND)
+    varfeat = varfeat.view(NP, -1, NK, ND)
+    keyvalue = keyvalue[:, None, :, :]
+    keyvalue = keyvalue + varfeat + query
+    keyvalue = torch.tanh(keyvalue)
+    keyvalue = keyvalue.view(-1, ND)
+
+    logit = masked_tensor(mask, 1).view(-1)
+    logit.addmv_(keyvalue, weight)
+    return logit.view(batch_size, -1)
+
+
+def masked_tensor(mask, heads):
+    size = list(mask.size())
+    size.insert(0, heads)
+    mask = mask[None].expand(size)
+    result = mask.new_zeros(size, dtype=torch.float32)
+    result[mask] = -math.inf
+    return result
+
+
+def vfaddmm(varfeat, mask, A, B):
+    varfeat = varfeat.permute(0, 2, 1)
+    return F.linear(varfeat, A.permute(1, 0), B)
+

greedrl/dense.py

@@ -0,0 +1,31 @@
+from torch import nn
+
+from .utils import get_act
+from .norm import Norm1D, Norm2D
+
+
+class Dense(nn.Module):
+
+    def __init__(self, input_dim, output_dim, bias=True, norm1d='none', norm2d='none', act='none'):
+        super(Dense, self).__init__()
+        assert norm1d == 'none' or norm2d == 'none', "one of [norm1d, norm2d] must be none"
+
+        if norm1d != 'none':
+            self.nn_norm = Norm1D(input_dim, norm1d)
+        elif norm2d != 'none':
+            self.nn_norm = Norm2D(input_dim, norm2d)
+        else:
+            self.nn_norm = None
+
+        self.nn_act = get_act(act)
+        self.nn_linear = nn.Linear(input_dim, output_dim, bias)
+
+    def weight(self):
+        return self.nn_linear.weight
+
+    def forward(self, x):
+        if self.nn_norm is not None:
+            x = self.nn_norm(x)
+        x = self.nn_act(x)
+        x = self.nn_linear(x)
+        return x

greedrl/encode.py

@@ -0,0 +1,349 @@
+import math
+import torch
+import torch.nn.functional as F
+
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from .norm import Norm1D, Norm2D
+from .dense import Dense
+from .utils import repeat
+from .feature import *
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, heads, hidden_dim):
+        super(MultiHeadAttention, self).__init__()
+        
+        assert hidden_dim % heads == 0
+
+        self.heads = heads
+        head_dim = hidden_dim // heads
+        self.alpha = 1 / math.sqrt(head_dim)
+
+        self.nn_Q = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+        self.nn_K = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+        self.nn_V = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+        self.nn_O = nn.Parameter(torch.Tensor(hidden_dim, hidden_dim))
+        
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, x, edge):
+        batch_size, item_num, hidden_dim = x.size()
+        size = (self.heads, batch_size, item_num, -1)
+        
+        x = x.reshape(-1, hidden_dim)
+        Q = torch.matmul(x, self.nn_Q).view(size)
+        K = torch.matmul(x, self.nn_K).view(size)
+        V = torch.matmul(x, self.nn_V).view(size)
+
+        heads_batch = self.heads * batch_size
+        Q = Q.view(heads_batch, item_num, -1)
+        K = K.view(heads_batch, item_num, -1).transpose(1, 2)
+        
+        if edge is not None:
+            S = edge.view(heads_batch, item_num, item_num)
+            S = S.baddbmm(Q, K, alpha=self.alpha)
+        else:
+            S = Q.new_zeros(heads_batch, item_num, item_num)
+            S = S.baddbmm_(Q, K, alpha=self.alpha)
+
+        S = S.view(self.heads, batch_size, item_num, item_num)
+
+        S = F.softmax(S, dim=-1)
+
+        x = torch.matmul(S, V).permute(1, 2, 0, 3)
+        x = x.reshape(batch_size, item_num, -1)
+        x = torch.matmul(x, self.nn_O)
+        return x
+
+
+class Encode(nn.Module):
+    def __init__(self, nn_args):
+        super(Encode, self).__init__()
+
+        self.nn_args = nn_args
+        self.worker_dim = nn_args['worker_dim']
+        self.task_dim = nn_args['task_dim']
+        self.edge_dim = nn_args['edge_dim']
+
+        self.embed_dict = nn_args['embed_dict']
+        self.feature_dict = nn_args['feature_dict']
+
+        layers = nn_args.setdefault('encode_layers', 3)
+        heads = nn_args.setdefault('encode_atten_heads', 8)
+        norm = nn_args.setdefault('encode_norm', 'instance')
+        hidden_dim = nn_args.setdefault('encode_hidden_dim', 128)
+        output_dim = nn_args.setdefault('decode_hidden_dim', 128)
+        output_heads = nn_args.setdefault('decode_atten_heads', 0)
+        
+        self.heads = heads
+        self.layers = layers
+
+        worker_dim = max(1, sum(self.worker_dim.values()))
+        task_dim = max(1, sum(self.task_dim.values()))
+
+        self.nn_dense_worker_start = Dense(worker_dim, hidden_dim)
+        self.nn_dense_worker_end = Dense(worker_dim, hidden_dim)
+        self.nn_dense_task = Dense(task_dim, hidden_dim)
+        
+        self.nn_norm_worker_task = Norm1D(hidden_dim, norm, True)
+
+        if len(self.edge_dim) > 0:
+            edge_dim = sum(self.edge_dim.values())
+            self.nn_dense_edge = Dense(edge_dim, heads)
+            self.nn_norm_edge = Norm2D(heads, norm, True)
+
+        nn_embed_dict = {}
+        for k, v in self.embed_dict.items():
+            nn_embed_dict[k] = nn.Embedding(v, hidden_dim)
+        self.nn_embed_dict = nn.ModuleDict(nn_embed_dict)
+
+        self.nn_attens = nn.ModuleList()
+        self.nn_denses = nn.ModuleList()
+        self.nn_norms1 = nn.ModuleList()
+        self.nn_norms2 = nn.ModuleList()
+        for i in range(layers):
+            self.nn_attens.append(MultiHeadAttention(heads, hidden_dim))
+            self.nn_denses.append(nn.Sequential(
+                                    Dense(hidden_dim, hidden_dim * 4),
+                                    Dense(hidden_dim * 4, hidden_dim, act='relu'),
+                                    ))
+            self.nn_norms1.append(Norm1D(hidden_dim, norm, True))
+            self.nn_norms2.append(Norm1D(hidden_dim, norm, True))
+
+        self.nn_finish = nn.Parameter(torch.Tensor(1, 1, hidden_dim))        
+        
+        if output_dim != hidden_dim:
+            self.nn_X = nn.Parameter(torch.Tensor(hidden_dim, output_dim))
+        else:
+            self.nn_X = None
+
+        if output_heads > 0:
+            assert output_dim % output_heads == 0
+            head_dim = output_dim // output_heads
+            self.nn_K = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+            self.nn_V = nn.Parameter(torch.Tensor(heads, hidden_dim, head_dim))
+        else:
+            self.nn_K = None
+            self.nn_V = None
+
+        for param in self.parameters():
+            stdv = 1 / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, problem, batch_size, worker_num, task_num, memopt=0):
+        worker_start, worker_end = self.encode_worker(problem, batch_size, worker_num)
+        task = self.encode_task(problem, batch_size, task_num)
+        X = torch.cat([worker_start, worker_end, task], 1)
+        X = self.nn_norm_worker_task(X)
+
+        if len(self.edge_dim) > 0:
+            edge = self.encode_edge(problem, batch_size, worker_num, task_num)
+            edge = self.nn_norm_edge(edge)
+            edge = edge.permute(3, 0, 1, 2).contiguous()
+        else:
+            edge = None
+        
+        #transformer encoding
+        for i in range(self.layers):
+            X = self.encode_layer(X, edge, i, memopt)
+
+        finish = repeat(self.nn_finish, X.size(0))
+        X = torch.cat([X, finish], 1)
+        if self.nn_X is not None:
+            X = torch.matmul(X, self.nn_X)
+        
+        if self.nn_K is not None:
+            batch_size, item_num, hidden_dim = X.size()
+            size = (self.heads, batch_size, item_num, -1)
+            X2 = X.reshape(-1, hidden_dim)
+            K = torch.matmul(X2, self.nn_K).view(size)
+            V = torch.matmul(X2, self.nn_V).view(size)
+        else:
+            K = torch.ones(0)
+            V = torch.ones(0)
+        return X, K, V
+    
+    def encode_layer(self, X, edge, i, memopt):
+        run_fn = self.encode_layer_fn(i, memopt)
+        if self.training and memopt > 6:
+            return checkpoint(run_fn, X, edge)
+        else:
+            return run_fn(X, edge)
+
+    def encode_layer_fn(self, i, memopt):
+        def run_fn(X, edge):
+            if self.training and memopt == 6:
+                X = X + checkpoint(self.nn_attens[i], X, edge)
+            else:
+                X = X + self.nn_attens[i](X, edge)
+            X = self.nn_norms1[i](X)
+
+            X = X + self.nn_denses[i](X)
+            X = self.nn_norms2[i](X)
+            return X
+
+        return run_fn
+    
+    def encode_worker(self, problem, batch_size, worker_num):
+        feature_list = []
+        for k, dim in self.worker_dim.items():
+            f = self.feature_dict.get(k)
+            if isinstance(f, GlobalCategory):
+                v = problem[f.name]
+                v = self.nn_embed_dict[k](v.long())
+            elif isinstance(f, ContinuousFeature):
+                v = problem[f.name]
+            else:
+                raise Exception("unsupported feature type: {}".format(type(f)))
+
+            if v.dim() == 2:
+                v = v[:, :, None]
+
+            assert dim == v.size(-1), \
+                "feature dim error, feature: {}, expected: {}, actual: {}".format(k, dim, v.size(-1))
+
+            feature_list.append(v.float())
+
+        if feature_list:
+            x = torch.cat(feature_list, 2)
+        else:
+            x = self.nn_finish.new_ones(batch_size, worker_num, 1)
+        return self.nn_dense_worker_start(x), self.nn_dense_worker_end(x)
+
+    def encode_task(self, problem, batch_size, task_num):
+        feature_list = []
+        for k, dim in self.task_dim.items():
+            f = self.feature_dict.get(k)
+            if isinstance(f, SparseLocalFeature):
+                v = problem[f.value]
+                assert v.dim() == 3, \
+                    "sparse local feature's dimension must 2, feature:{}".format(k)
+                v = v.clamp(0, 1).sum(2, dtype=v.dtype)
+            elif isinstance(f, GlobalCategory):
+                v = problem[f.name]
+                v = self.nn_embed_dict[k](v.long())
+            elif isinstance(f, LocalFeature):
+                v = problem[f.name]
+                assert v.dim() == 3, \
+                    "local feature's dimension must 2, feature:{}".format(k)
+                v = v.clamp(0, 1).sum(2, dtype=v.dtype)
+            elif isinstance(f, ContinuousFeature):
+                v = problem[f.name]
+            else:
+                raise Exception("unsupported feature type: {}".format(type(f)))
+
+            if v.dim() == 2:
+                v = v[:, :, None]
+
+            assert dim == v.size(-1), \
+                "feature dim error, feature: {}, expected: {}, actual: {}".format(k, dim, v.size(-1))
+
+            feature_list.append(v.float())
+
+        if feature_list:
+            x = torch.cat(feature_list, 2)
+        else:
+            x = self.nn_finish.new_ones(batch_size, task_num, 1)
+        return self.nn_dense_task(x)
+
+    def encode_edge(self, problem, batch_size, worker_num, task_num):
+        NP = batch_size
+        NW = worker_num
+        NT = task_num
+        NWW = NW + NW
+        feature_list = []
+        for k, dim in self.edge_dim.items():
+            f = self.feature_dict.get(k)
+            if isinstance(f, LocalCategory):
+                assert f.name.startswith("task_")
+
+                v = problem[k]
+                v1 = v[:, :, None]
+                v2 = v[:, None, :]
+
+                v = torch.zeros(NP, NWW + NT, NWW + NT,
+                                dtype=v.dtype, device=v.device)
+                v[:, NWW:, NWW:] = ((v1 == v2) & (v1 >= 0))
+            elif isinstance(f, LocalFeature):
+                assert f.name.startswith("task_")
+
+                v = problem[k].float()
+                dot_product = torch.matmul(v, v.transpose(-1, -2))
+                v_norm = v.norm(dim=2) + 1e-10
+                v1_norm = v_norm[:, :, None]
+                v2_norm = v_norm[:, None, :]
+
+                v = torch.zeros(NP, NWW + NT, NWW + NT,
+                                dtype=v.dtype, device=v.device)
+                v[:, NWW:, NWW:] = dot_product / v1_norm / v2_norm
+            elif isinstance(f, SparseLocalFeature):
+                assert NP == 1
+                assert f.index.startswith("task_")
+                assert f.value.startswith("task_")
+
+                index = problem[f.index]
+                value = problem[f.value].float()
+                
+                NV = index.max().item() + 1
+                spv = value.reshape(-1).tolist()
+                spi = index.reshape(-1).tolist()
+
+                device = value.device
+                spj = torch.arange(NT, device=device)
+                spj = spj[:, None].expand_as(index)
+                spj = spj.reshape(-1).tolist()
+                
+                value1 = torch.sparse_coo_tensor([spj, spi], spv, (NT, NV), device=device)
+                value2 = torch.sparse_coo_tensor([spi, spj], spv, (NV, NT), device=device)
+                
+                value1 = value1.coalesce()
+                value2 = value2.coalesce()
+                cosine = torch.sparse.mm(value1, value2).to_dense()
+                
+                norm = value.norm(dim=-1).reshape(-1)
+                norm1 = norm[:, None].expand(-1, NT)
+                norm2 = norm[None, :].expand(NT, -1)
+                cosine = cosine / (norm1 * norm2 + 1e-10)
+
+                v = torch.zeros(NP, NWW + NT, NWW + NT,
+                                dtype=value.dtype, device=value.device)
+                v[:, NWW:, NWW:] = cosine
+
+            elif isinstance(f, ContinuousFeature):
+                if f.name.endswith("_matrix"):
+                    v = problem[k]
+                elif f.name.startswith("worker_task_"):
+                    v = problem[k]
+                    if v.dim() == 3:
+                        new_v = torch.zeros(NP, NWW + NT, NWW + NT,
+                                            dtype=v.dtype, device=v.device)
+                    else:
+                        new_v = torch.zeros(NP, NWW + NT, NWW + NT, v.size(3),
+                                            dtype=v.dtype, device=v.device)
+                    problem_index = torch.arange(NP, device=v.device)[:, None, None]
+                    worker_index = torch.arange(NW, device=v.device)[None, :, None]
+                    task_index = torch.arange(NT, device=v.device)[None, None, :] + NW + NW
+                    new_v[problem_index, worker_index, task_index] = v
+                    new_v[problem_index, task_index, worker_index] = v
+                    new_v[problem_index, worker_index + NW, task_index] = v
+                    new_v[problem_index, task_index, worker_index + NW] = v
+                    v = new_v
+                else:
+                    raise Exception("feature: {}".format(f.name))
+            else:
+                raise Exception("feature: {}, type: {}".format(k, type(f)))
+
+            if v.dim() == 3:
+                v = v[:, :, :, None]
+
+            assert dim == v.size(-1), \
+                "feature: {}, expected: {}, actual: {}".format(k, dim, v.size(-1))
+
+            feature_list.append(v.float())
+
+        x = torch.cat(feature_list, 3)
+        return self.nn_dense_edge(x)
+

greedrl/feature.py

@@ -0,0 +1,63 @@
+def continuous_feature(name):
+    return ContinuousFeature(name)
+
+
+class ContinuousFeature:
+    def __init__(self, name):
+        self.name = name
+
+
+def global_category(name, size):
+    return GlobalCategory(name, size)
+
+
+class GlobalCategory:
+    def __init__(self, name, size):
+        self.name = name
+        self.size = size
+
+
+def local_category(name):
+    return LocalCategory(name)
+
+
+class LocalCategory:
+    def __init__(self, name):
+        assert name.startswith('task_'), \
+            "only task feature supported: {}".format(name)
+        self.name = name
+
+
+def local_feature(name):
+    return LocalFeature(name)
+
+
+class LocalFeature:
+    def __init__(self, name):
+        assert name.startswith('task_'), \
+            "only task feature supported: {}".format(name)
+        self.name = name
+
+
+def sparse_local_feature(index, value):
+    return SparseLocalFeature(index, value)
+
+
+class SparseLocalFeature:
+    def __init__(self, index, value):
+        assert index.startswith('task_'), \
+            "only task feature supported for index: {}".format(index)
+        assert value.startswith('task_'), \
+            "only task feature supported for value: {}".format(value)
+
+        self.index = index
+        self.value = value
+
+
+def variable_feature(name):
+    return VariableFeature(name)
+
+
+class VariableFeature:
+    def __init__(self, name):
+        self.name = name

greedrl/function.py

@@ -0,0 +1,5 @@
+import greedrl_c
+from greedrl_c import task_group_priority
+from greedrl_c import task_group_split
+
+

greedrl/norm.py

@@ -0,0 +1,25 @@
+import torch
+
+from torch import nn
+
+
+class Norm1D(nn.Module):
+
+    def __init__(self, dim, ntype='batch', affine=False):
+        super(Norm1D, self).__init__()
+        clazz_dict = {'batch': nn.BatchNorm1d, 'instance': nn.InstanceNorm1d}
+        self.nn_norm = clazz_dict[ntype](dim, eps=1e-10, affine=affine)
+
+    def forward(self, x):
+        return self.nn_norm(x.permute(0, 2, 1)).permute(0, 2, 1)
+
+
+class Norm2D(nn.Module):
+
+    def __init__(self, dim, ntype='batch', affine=False):
+        super(Norm2D, self).__init__()
+        clazz_dict = {'batch': nn.BatchNorm2d, 'instance': nn.InstanceNorm2d}
+        self.nn_norm = clazz_dict[ntype](dim, eps=1e-10, affine=affine)
+
+    def forward(self, x):
+        return self.nn_norm(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

greedrl/pyenv.py

@@ -0,0 +1,383 @@
+import torch
+import json
+import math
+
+from collections import OrderedDict
+from .const import *
+from .utils import to_list
+from .norm import Norm1D, Norm2D
+from .variable import AttributeVariable, WorkerTaskSequence
+
+
+class PyEnv(object):
+
+    def __init__(self, problem, batch_size, sample_num, nn_args):
+        super(PyEnv, self).__init__()
+
+        self._problem = problem
+        self._batch_size = batch_size
+        self._sample_num = sample_num
+        self._debug = -1
+
+        self._NW = problem.worker_num
+        self._NWW = problem.worker_num * 2
+        self._NT = problem.task_num
+        self._NWWT = self._NWW + self._NT
+
+        self._feats_dict = nn_args['feature_dict']
+        self._vars_dim = nn_args['variable_dim']
+
+        self._vars_dict = {}
+        self._vars = [var(problem, batch_size, sample_num) for var in problem.variables]
+        for variable in self._vars:
+            save_variable_version(variable)
+            assert variable.name not in self._vars_dict, \
+                "duplicated variable, name: {}".format(variable.name)
+            self._vars_dict[variable.name] = variable
+
+        self._constraint = problem.constraint()
+        self._objective = problem.objective()
+
+        self._worker_index = torch.full((self._batch_size,), -1,
+                                        dtype=torch.int64,
+                                        device=problem.device)
+
+        self._batch_index = torch.arange(self._batch_size,
+                                         dtype=torch.int64,
+                                         device=problem.device)
+
+        self._problem_index = torch.div(self._batch_index, sample_num, rounding_mode='trunc') #  self._batch_index // sample_num
+
+        self._feasible = torch.ones(self._batch_size,
+                                    dtype=torch.bool,
+                                    device=problem.device)
+
+        self._cost = torch.zeros(self._batch_size, self._NT * 2,
+                                 dtype=torch.float32,
+                                 device=problem.device)
+
+        self._mask = torch.zeros(self._batch_size,
+                                 self._NWWT + 1,
+                                 dtype=torch.bool,
+                                 device=problem.device)
+
+        self._worker_task_sequence = torch.full((self._batch_size, self._NT * 2, 3), -1,
+                                                dtype=torch.int64,
+                                                device=problem.device)
+        self._step = 0
+        self.register_variables(self._constraint)
+        self._finished = self._constraint.finished()
+
+        if hasattr(self._constraint, 'mask_worker_start'):
+            self.register_variables(self._constraint)
+            mask_start = self._constraint.mask_worker_start()
+        else:
+            mask_start = False
+
+        self._mask[:, :self._NW] = mask_start
+        self._mask[:, self._NW:] = True
+
+        if self._debug >= 0:
+            print("\n$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
+            print("new env")
+            print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$\n")
+
+    def time(self):
+        return self._step
+
+    def step(self, chosen):
+        with torch.no_grad():
+            self._do_step(chosen)
+
+    def _do_step(self, chosen):
+        if self._debug >= 0:
+            print("----------------------------------------------------------------------")
+            feasible = self._feasible & ~self._mask[self._problem_index, chosen]
+            print("feasible={}".format(feasible[self._debug].tolist()))
+
+        is_start = (chosen >= 0) & (chosen < self._NW)
+        if torch.any(is_start):
+            b_index = self._batch_index[is_start]
+            p_index = self._problem_index[is_start]
+            w_index = chosen[is_start]
+            self.step_worker_start(b_index, p_index, w_index)
+
+        is_end = (chosen >= self._NW) & (chosen < self._NWW)
+        if torch.any(is_end):
+            b_index = self._batch_index[is_end]
+            p_index = self._problem_index[is_end]
+            w_index = chosen[is_end] - self._NW
+            self.step_worker_end(b_index, p_index, w_index)
+
+        is_task = (chosen >= self._NWW) & (chosen < self._NWWT)
+        if torch.any(is_task):
+            b_index = self._batch_index[is_task]
+            p_index = self._problem_index[is_task]
+            t_index = chosen[is_task] - self._NWW
+            step_task_b_index = b_index
+            self.step_task(b_index, p_index, t_index)
+        else:
+            step_task_b_index = None
+
+        is_finish = chosen == self._NWWT
+        if torch.any(is_finish):
+            b_index = self._batch_index[is_finish]
+            self._worker_task_sequence[b_index, self._step, 0] = GRL_FINISH
+            self._worker_task_sequence[b_index, self._step, 1] = 0
+            self._worker_task_sequence[b_index, self._step, 2] = -1
+
+        self.update_mask(step_task_b_index)
+
+        for var in self._vars:
+            check_variable_version(var)
+
+        if self._debug >= 0:
+            print("worker_task_sequence[{}]={}".format(self._step,
+                                                       self._worker_task_sequence[self._debug, self._step].tolist()))
+            for var in self._vars:
+                if var.value is None:
+                    print("{}={}".format(var.name, None))
+                elif isinstance(var, AttributeVariable):
+                    print("{}={}".format(var.name, to_list(var.value)))
+                else:
+                    print("{}={}".format(var.name, to_list(var.value[self._debug])))
+
+        self._step += 1
+        if self._step >= self._cost.size(1):
+            cost = torch.zeros(self._batch_size, self._step + self._NT,
+                               dtype=torch.float32,
+                               device=chosen.device)
+            cost[:, 0:self._step] = self._cost;
+            self._cost = cost
+
+            worker_task_sequence = torch.full((self._batch_size, self._step + self._NT, 3), -1,
+                                              dtype=torch.int64,
+                                              device=chosen.device)
+            worker_task_sequence[:, 0:self._step, :] = self._worker_task_sequence
+            self._worker_task_sequence = worker_task_sequence
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        self._worker_task_sequence[b_index, self._step, 0] = GRL_WORKER_START
+        self._worker_task_sequence[b_index, self._step, 1] = w_index
+        self._worker_task_sequence[b_index, self._step, 2] = -1
+        for var in self._vars:
+            if hasattr(var, 'step_worker_start'):
+                var.step_worker_start(b_index, p_index, w_index)
+                save_variable_version(var)
+
+        if hasattr(self._objective, 'step_worker_start'):
+            self.register_variables(self._objective, b_index)
+            self.update_cost(self._objective.step_worker_start(), b_index)
+
+        self._worker_index[b_index] = w_index
+        self._mask[b_index, :self._NWW] = True
+        self._mask[b_index, self._NWW:] = False
+
+    def step_worker_end(self, b_index, p_index, w_index):
+        self._worker_task_sequence[b_index, self._step, 0] = GRL_WORKER_END
+        self._worker_task_sequence[b_index, self._step, 1] = w_index
+        self._worker_task_sequence[b_index, self._step, 2] = -1;
+
+        for var in self._vars:
+            if hasattr(var, 'step_worker_end'):
+                var.step_worker_end(b_index, p_index, w_index)
+                save_variable_version(var)
+
+        if hasattr(self._objective, 'step_worker_end'):
+            self.register_variables(self._objective, b_index)
+            self.update_cost(self._objective.step_worker_end(), b_index)
+
+        self._worker_index[b_index] = -1
+
+        self.register_variables(self._constraint, b_index)
+        self._finished[b_index] |= self._constraint.finished()
+        if hasattr(self._constraint, 'mask_worker_start'):
+            mask_start = self._constraint.mask_worker_start()
+        else:
+            mask_start = False
+
+        self._mask[b_index, :self._NW] = mask_start
+        self._mask[b_index, self._NW:] = True
+
+    def step_task(self, b_index, p_index, t_index):
+        self._worker_task_sequence[b_index, self._step, 0] = GRL_TASK
+        self._worker_task_sequence[b_index, self._step, 1] = t_index
+
+        for var in self._vars:
+            if not hasattr(var, 'step_task'):
+                continue
+            elif var.step_task.__code__.co_argcount == 4:
+                var.step_task(b_index, p_index, t_index)
+            else:
+                var.step_task(b_index, p_index, t_index, None)
+            save_variable_version(var)
+
+        if hasattr(self._constraint, 'do_task'):
+            self.register_variables(self._constraint, b_index)
+            done = self._constraint.do_task()
+            self._worker_task_sequence[b_index, self._step, 2] = done.long()
+
+            for var in self._vars:
+                if not hasattr(var, 'step_task'):
+                    continue
+                elif var.step_task.__code__.co_argcount == 4:
+                    pass
+                else:
+                    check_variable_version(var)
+                    var.step_task(b_index, p_index, t_index, done)
+                    save_variable_version(var)
+        else:
+            done = None
+
+        if hasattr(self._objective, 'step_task'):
+            self.register_variables(self._objective, b_index)
+            self.update_cost(self._objective.step_task(), b_index)
+
+        if hasattr(self._constraint, 'mask_worker_end'):
+            self.register_variables(self._constraint, b_index)
+            mask_end = self._constraint.mask_worker_end()
+        else:
+            mask_end = False
+
+        w_index = self._NW + self._worker_index[b_index]
+        self._mask[b_index, w_index] = mask_end
+        self._mask[b_index, self._NWW:] = False
+        return done
+
+    def update_cost(self, cost, b_index=None):
+        if isinstance(cost, tuple):
+            cost, feasible = cost
+            if b_index is None:
+                self._feasible &= feasible
+            else:
+                self._feasible[b_index] &= feasible
+
+        if isinstance(cost, torch.Tensor):
+            cost = cost.float()
+        else:
+            assert type(cost) in (int, float), "unexpected cost's type: {}".format(type(cost))
+
+        if b_index is None:
+            self._cost[:, self._step] = cost
+        else:
+            self._cost[b_index, self._step] = cost
+
+    def update_mask(self, step_task_b_index):
+        self._mask |= self._finished[:, None]
+        self._mask[:, -1] = ~self._finished
+        self.register_variables(self._constraint)
+        self._mask[:, self._NWW:self._NWWT] |= self._constraint.mask_task()
+
+        if step_task_b_index is not None:
+            b_index = step_task_b_index
+            w_index = self._NW + self._worker_index[b_index]
+            task_mask = self._mask[b_index, self._NWW:self._NWWT]
+            self._mask[b_index, w_index] &= ~torch.all(task_mask, 1)
+
+    def batch_size():
+        return self._batch_size
+
+    def sample_num():
+        return self._sample_num
+
+    def mask(self):
+        return self._mask.clone()
+
+    def cost(self):
+        return self._cost[:, 0:self._step]
+
+    def feasible(self):
+        return self._feasible
+
+    def worker_task_sequence(self):
+        return self._worker_task_sequence[:, 0:self._step]
+
+    def var(self, name):
+        return self._vars_dict[name].value
+
+    def register_variables(self, obj, b_index=None, finished=False):
+        for var in self._vars:
+            if var.value is None or b_index is None \
+                    or isinstance(var, AttributeVariable):
+                value = var.value
+            else:
+                value = var.value[b_index]
+            obj.__dict__[var.name] = value
+
+            if not hasattr(var, 'ext_values'):
+                continue
+
+            for k, v in var.ext_values.items():
+                k = var.name + '_' + k
+                obj.__dict__[k] = v[b_index]
+
+    def finished(self):
+        return self._finished
+
+    def all_finished(self):
+        return torch.all(self.finished())
+
+    def finalize(self):
+        self._worker_task_sequence[:, self._step, 0] = GRL_FINISH
+        self._worker_task_sequence[:, self._step, 1] = 0
+        self._worker_task_sequence[:, self._step, 2] = -1
+
+        for var in self._vars:
+            if hasattr(var, 'step_finish'):
+                var.step_finish(self.worker_task_sequence())
+
+        if hasattr(self._objective, 'step_finish'):
+            self.register_variables(self._objective, finished=True)
+            self.update_cost(self._objective.step_finish())
+
+        self._step += 1
+
+    def make_feat(self):
+        with torch.no_grad():
+            return self.do_make_feat()
+
+    def do_make_feat(self):
+        if not self._vars_dim:
+            return None
+
+        feature_list = []
+        for k, dim in self._vars_dim.items():
+            f = self._feats_dict[k]
+            var = self._vars_dict[f.name]
+            v = var.make_feat()
+            if v.dim() == 2:
+                v = v[:, :, None]
+
+            assert dim == v.size(-1), \
+                "feature dim error, feature: {}, expected: {}, actual: {}".format(k, dim, v.size(-1))
+            feature_list.append(v.float())
+
+        v = torch.cat(feature_list, 2)
+        u = v.new_zeros(v.size(0), self._NWW, v.size(2))
+        f = v.new_zeros(v.size(0), 1, v.size(2))
+        v = torch.cat([u, v, f], 1).permute(0, 2, 1)
+
+        v[self._mask[:, None, :].expand(v.size())] = 0
+
+        norm = v.new_ones(self._mask.size())
+        norm[self._mask] = 0
+        norm = norm.sum(1) + 1e-10
+        norm = norm[:, None, None]
+
+        avg = v.sum(-1, keepdim=True) / norm
+        v = v - avg
+
+        std = v.norm(dim=-1, keepdim=True) / norm + 1e-10
+        v = v / std
+        return v.contiguous()
+
+
+def save_variable_version(var):
+    if isinstance(var.value, torch.Tensor):
+        var.__version__ = var.value._version
+
+
+def check_variable_version(var):
+    if isinstance(var.value, torch.Tensor):
+        assert var.__version__ == var.value._version, \
+            "variable's value is modified, name: {}".format(var.name)

greedrl/solver.py

@@ -0,0 +1,625 @@
+import os
+import sys
+import math
+import copy
+import time
+import queue
+import inspect
+import torch
+import numpy as np
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from .agent import Agent, parse_nn_args
+from .utils import repeat, get_default_device, cutime_stats
+from .variable import TaskDemandNow
+
+from torch.nn.utils import clip_grad_norm_, parameters_to_vector, vector_to_parameters
+from torch.utils.data import Dataset, IterableDataset, DataLoader
+from torch.optim.lr_scheduler import MultiStepLR
+
+
+class Problem(object):
+    def __init__(self, isbatch=False):
+        self.isbatch = isbatch
+        self.features = []
+        self.environment = None
+
+    def pin_memory(self):
+        for k, v in self.feats.items():
+            self.feats[k] = v.pin_memory()
+        return self
+
+    def __getattr__(self, name):
+        if name not in ('solution'):
+            raise AttributeError()
+        return self.feats.get(name)
+
+
+class Solution(object):
+    def __init__(self, cost=None):
+        self.cost = cost
+        self.worker_task_sequence = None
+
+
+class WrapDataset(Dataset):
+    def __init__(self, dataset, solver):
+        self._dataset = [solver.to_batch(p) for p in dataset]
+
+    def __getitem__(self, index):
+        return self._dataset[index]
+
+    def __len__(self):
+        return len(self._dataset)
+
+
+class WrapIterator:
+    def __init__(self, iterator, solver):
+        self._iterator = iterator
+        self._solver = solver
+
+    def __next__(self):
+        p = next(self._iterator)
+        p = self._solver.to_batch(p, False)
+        return p
+
+
+class WrapIterableDataset(IterableDataset):
+    def __init__(self, dataset, solver):
+        self._dataset = dataset
+        self._solver = solver
+
+    def __iter__(self):
+        return WrapIterator(iter(self._dataset), self._solver)
+
+
+class CyclicIterator:
+    def __init__(self, iterable):
+        self._iterable = iterable
+        self._iterator = iter(iterable)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._iterator = iter(self._iterable)
+            return next(self._iterator)
+
+
+class BufferedIterator:
+    def __init__(self, iterator, size, reuse):
+        self._iterator = iterator
+        self._reuse = reuse
+        self._queue = queue.Queue(size)
+        self._buffer = []
+        self._iter_step = 0
+
+    def __next__(self):
+        if not self._queue.full() or self._iter_step % self._reuse == 0:
+            problem = next(self._iterator)
+            if self._queue.full():
+                index = self._queue.get()
+                self._buffer[index] = problem
+            else:
+                index = len(self._buffer)
+                self._buffer.append(problem)
+            self._queue.put(index)
+        self._iter_step += 1
+        index = torch.randint(0, len(self._buffer), (1,)).item()
+        return self._buffer[index]
+
+
+class Solver(object):
+    def __init__(self, device=None, nn_args=None):
+
+        if device is None:
+            self.device = get_default_device()
+        elif device == 'cuda':
+            self.device = get_default_device()
+            assert self.device.type == 'cuda', 'no cuda device available!'
+        else:
+            self.device = torch.device(device)
+
+        if nn_args is None:
+            nn_args = {}
+        self.nn_args = nn_args
+
+        self.agent = None
+
+    def parse_nn_args(self, problem):
+        parse_nn_args(problem, self.nn_args)
+
+    def new_agent(self):
+        return Agent(self.nn_args)
+
+    def train(self, agent_filename, train_dataset, valid_dataset, **kwargs):
+        if dist.is_initialized():
+            torch.manual_seed(torch.initial_seed() + dist.get_rank() * 20000)
+
+        train_dataset_workers = kwargs.pop('train_dataset_workers', 1)
+        train_dataset_buffers = kwargs.pop('train_dataset_buffers', 2)
+        valid_dataset_workers = kwargs.pop('valid_dataset_workers', 1)
+        valid_dataset_buffers = kwargs.pop('valid_dataset_buffers', 2)
+
+        train_dataset = self.wrap_dataset(train_dataset, train_dataset_workers,
+                                          train_dataset_buffers, torch.initial_seed() + 1)
+        valid_dataset = self.wrap_dataset(valid_dataset, valid_dataset_workers,
+                                          valid_dataset_buffers, torch.initial_seed() + 10001)
+
+        if self.device.type == 'cuda':
+            with torch.cuda.device(cuda_or_none(self.device)):
+                self.do_train(agent_filename, train_dataset, valid_dataset, **kwargs)
+        else:
+            self.do_train(agent_filename, train_dataset, valid_dataset, **kwargs)
+
+    def do_train(self, agent_filename, train_dataset, valid_dataset, reuse_buffer=0, reuse_times=1, on_policy=True,
+                 advpow=1, batch_size=512, topk_size=1, init_lr=0.0001, sched_lr=(int(1e10),), gamma_lr=0.5,
+                 warmup_steps=100, log_steps=-1, optim_steps=1, valid_steps=100, max_steps=int(1e10), memopt=1):
+
+        for arg in inspect.getfullargspec(self.do_train)[0][1:]:
+            if arg not in ('train_dataset', 'valid_dataset'):
+                print("train_args: {} = {}".format(arg, locals()[arg]))
+
+        if log_steps < 0:
+            log_steps = valid_steps
+
+        train_dataset = CyclicIterator(train_dataset)
+        if reuse_buffer > 0:
+            train_dataset = BufferedIterator(train_dataset, reuse_buffer, reuse_times)
+
+        valid_dataset = list(valid_dataset)
+
+        if dist.is_initialized() and dist.get_rank() != 0:
+            dist.barrier()
+
+        if agent_filename is not None and os.path.exists(agent_filename):
+            saved_state = torch.load(agent_filename, map_location='cpu')
+            self.nn_args = saved_state['nn_args']
+        else:
+            saved_state = None
+            self.parse_nn_args(valid_dataset[0])
+
+        step = 0
+        start_step = 0
+        self.agent = self.new_agent().train()
+        self.agent.to(self.device)
+        self.print_nn_args()
+
+        best_agent = copy.deepcopy(self.agent).eval()
+        min_valid_cost = math.inf
+
+        optimizer = torch.optim.Adam(self.agent.parameters(), lr=init_lr)
+        scheduler = MultiStepLR(optimizer, milestones=sched_lr, gamma=gamma_lr)
+
+        def do_save_state(rng_state, cuda_rng_state):
+            if agent_filename is not None:
+                save_data = {'step': step, 'rng_state': rng_state}
+                if cuda_rng_state is not None:
+                    save_data['cuda_rng_state'] = cuda_rng_state
+                save_data['nn_args'] = self.agent.nn_args_dict()
+                save_data['agent_state'] = self.agent.state_dict()
+                save_data['best_agent_state'] = best_agent.state_dict()
+                save_data['optimizer_state'] = optimizer.state_dict()
+                save_data['scheduler_state'] = scheduler.state_dict()
+                torch.save(save_data, agent_filename)
+
+        def valid_sched_save(step):
+            if dist.is_initialized():
+                params = parameters_to_vector(self.agent.parameters())
+                params_clone = params.clone()
+                dist.broadcast(params_clone, 0)
+                assert torch.all(params == params_clone)
+
+            rng_state = torch.get_rng_state()
+            cuda_rng_state = None
+            if self.device.type == 'cuda':
+                cuda_rng_state = torch.cuda.get_rng_state(self.device)
+
+            print("{} - step={}, validate...".format(time.strftime("%Y-%m-%d %H:%M:%S"), step))
+            sys.stdout.flush()
+
+            if self.device.type == 'cuda':
+                torch.cuda.synchronize(self.device)
+            start_time = time.time()
+            valid_result = self.validate(valid_dataset, batch_size)
+            avg_cost1, avg_cost2, avg_feasible = valid_result
+            if self.device.type == 'cuda':
+                torch.cuda.synchronize(self.device)
+
+            duration = time.time() - start_time
+
+            if step > 0:
+                scheduler.step()
+
+            if not dist.is_initialized() or dist.get_rank() == 0:
+                do_save_state(rng_state, cuda_rng_state)
+
+            strftime = time.strftime("%Y-%m-%d %H:%M:%S")
+            print("{} - step={}, cost=[{:.6g}, {:.6g}], feasible={:.0%}".format(
+                strftime, step, avg_cost1, avg_cost2, avg_feasible))
+            print("{} - step={}, min_valid_cost={:.6g}, time={:.3f}s".format(
+                strftime, step, min(min_valid_cost, avg_cost2), duration))
+            print("---------------------------------------------------------------------------------------")
+            sys.stdout.flush()
+            return avg_cost2
+
+        if saved_state is not None:
+            start_step = saved_state['step']
+
+            if not dist.is_initialized() or dist.get_rank() == 0:
+                torch.set_rng_state(saved_state['rng_state'])
+                if torch.cuda.is_available():
+                    torch.cuda.set_rng_state(saved_state['cuda_rng_state'], self.device)
+
+            best_agent.load_state_dict(saved_state['best_agent_state'])
+            self.agent.load_state_dict(saved_state['best_agent_state'])
+
+            # if 'agent_state' in saved_state:
+            #    self.agent.load_state_dict(saved_state['agent_state'])
+            # else:
+            #    self.agent.load_state_dict(saved_state['best_agent_state'])
+
+            if 'optimizer_state' in saved_state:
+                optimizer.load_state_dict(saved_state['optimizer_state'])
+            if 'scheduler_state' in saved_state:
+                scheduler.load_state_dict(saved_state['scheduler_state'])
+        else:
+            if dist.is_initialized() and dist.get_rank() == 0:
+                rng_state = torch.get_rng_state()
+                cuda_rng_state = None
+                if self.device.type == 'cuda':
+                    cuda_rng_state = torch.cuda.get_rng_state(self.device)
+                do_save_state(rng_state, cuda_rng_state)
+
+        if dist.is_initialized() and dist.get_rank() == 0:
+            dist.barrier()
+
+        for step in range(start_step, max_steps):
+            if step % valid_steps == 0:
+                valid_cost = valid_sched_save(step)
+                if valid_cost < min_valid_cost:
+                    best_agent.load_state_dict(self.agent.state_dict())
+                    min_valid_cost = valid_cost
+
+            start_time = time.time()
+
+            # problem
+            with torch.no_grad():
+                problem = next(train_dataset)
+                if step < warmup_steps:
+                    batch_size_now = batch_size // 2
+                else:
+                    batch_size_now = batch_size
+                problem = self.to_device(problem)
+
+            if not on_policy:
+                data_agent = best_agent
+            else:
+                data_agent = self.agent
+
+            data_agent.eval()
+
+            # solution
+            if topk_size > 1:
+                with torch.no_grad():
+                    batch_size_topk = batch_size_now * topk_size
+                    env, logp = data_agent(problem, batch_size_topk)
+                    cost = env.cost().sum(1).float()
+                    solution = env.worker_task_sequence()
+
+                    NP = problem.batch_size
+                    NK = batch_size_now // NP
+                    NS = solution.size(1)
+
+                    cost = cost.view(NP, -1)
+                    cost, kidx = cost.topk(NK, 1, False, False)
+                    cost = cost.view(-1)
+                    kidx = kidx[:, :, None, None].expand(-1, -1, NS, 3)
+                    solution = solution.view(NP, -1, NS, 3)
+                    solution = solution.gather(1, kidx).view(-1, NS, 3)
+
+            elif not on_policy:
+                with torch.no_grad():
+                    env, logp = data_agent(problem, batch_size_now)
+                    cost = env.cost().sum(1).float()
+                    solution = env.worker_task_sequence()
+            else:
+                self.agent.train()
+                env, logp = self.agent(problem, batch_size_now, memopt=memopt)
+                cost = env.cost().sum(1).float()
+                solution = env.worker_task_sequence()
+
+            self.agent.train()
+
+            # advantage
+            with torch.no_grad():
+                NP = problem.batch_size
+                if topk_size > 1:
+                    baseline = cost.view(NP, -1).max(1)[0]
+                else:
+                    baseline = cost.view(NP, -1).mean(1)
+                baseline = repeat(baseline, cost.size(0) // NP)
+                adv = (cost - baseline)[:, None]
+                adv_norm = adv.norm()
+                if adv_norm > 0:
+                    adv = adv / adv.norm() * adv.size(0)
+                    adv = adv.sign() * adv.abs().pow(advpow)
+
+            # backward
+            if topk_size > 1 or not on_policy:
+                env, logp = self.agent(problem, batch_size_now, solution=solution, memopt=memopt)
+
+            loss = adv * logp
+            loss = loss.mean()
+            loss.backward()
+
+            if step % optim_steps == 0:
+                if dist.is_initialized():
+                    params = filter(lambda a: a.grad is not None, self.agent.parameters())
+                    grad_list = [param.grad for param in params]
+                    grad_vector = parameters_to_vector(grad_list)
+                    dist.all_reduce(grad_vector, op=dist.ReduceOp.SUM)
+                    vector_to_parameters(grad_vector, grad_list)
+
+                grad_norm = clip_grad_norm_(self.agent.parameters(), 1)
+                optimizer.step()
+                optimizer.zero_grad()
+
+            if step % log_steps == 0:
+                strftime = time.strftime("%Y-%m-%d %H:%M:%S")
+                lr = optimizer.param_groups[0]['lr']
+                duration = time.time() - start_time
+                with torch.no_grad():
+                    p = logp.to(torch.float64).sum(1).exp().mean()
+                print("{} - step={}, grad={:.6g}, lr={:.6g}, p={:.6g}".format(
+                    strftime, step, grad_norm, lr, p))
+
+                print("{} - step={}, cost={:.6g}, time={:.3f}s".format(strftime, step, cost.mean(), duration))
+                print("---------------------------------------------------------------------------------------")
+                sys.stdout.flush()
+
+        valid_sched_save(step)
+
+    def solve(self, problem, greedy=False, batch_size=512):
+        if self.device.type == 'cuda':
+            with torch.cuda.device(cuda_or_none(self.device)):
+                return self.do_solve(problem, greedy, batch_size)
+        else:
+            return self.do_solve(problem, greedy, batch_size)
+
+    def do_solve(self, problem, greedy, batch_size):
+        isbatch = problem.isbatch
+        problem = self.to_batch(problem)
+        problem = self.to_device(problem)
+
+        if self.agent is None:
+            self.parse_nn_args(problem)
+            self.agent = self.new_agent()
+            self.agent.to(self.device)
+
+        self.agent.eval()
+
+        with torch.no_grad():
+            env, prob = self.agent(problem, batch_size, greedy, problem.solution)
+
+        NP = problem.batch_size
+        NR = prob.size(0) // NP
+
+        prob = prob.view(NP, NR, -1)
+        cost = env.cost().sum(1).view(NP, NR)
+        feasible = env.feasible().view(NP, NR)
+        size = list(env.worker_task_sequence().size())
+        size = [NP, NR] + size[1:]
+        worker_task_sequence = env.worker_task_sequence().view(size)
+
+        p_index = torch.arange(NP)
+        base_cost = cost.max() + 1
+        cost[~feasible] += base_cost
+        cost, s_index = cost.min(1)
+        feasible = feasible[p_index, s_index]
+        cost[~feasible] -= base_cost
+        probability = prob[p_index, s_index].exp()
+        worker_task_sequence = worker_task_sequence[p_index, s_index]
+
+        if isbatch:
+            solution = Solution(cost)
+            solution.feasible = feasible
+            solution.probability = probability
+            solution.worker_task_sequence = worker_task_sequence
+        else:
+            solution = Solution(cost.item())
+            solution.feasible = feasible.item()
+            solution.probability = probability.squeeze(0)
+            solution.worker_task_sequence = worker_task_sequence.squeeze(0)
+
+        return solution
+
+    def load_agent(self, filename, strict=True):
+        if self.device.type == 'cuda':
+            with torch.cuda.device(cuda_or_none(self.device)):
+                self.do_load_agent(filename, strict)
+        else:
+            self.do_load_agent(filename, strict)
+
+    def do_load_agent(self, filename, strict=True):
+        saved_state = torch.load(filename, map_location='cpu')
+        self.nn_args = saved_state['nn_args']
+
+        self.agent = self.new_agent()
+        self.agent.to(self.device)
+        self.agent.load_state_dict(saved_state['best_agent_state'], strict)
+        self.print_nn_args()
+
+    def to_batch(self, problem, pin_memory=True):
+        assert not hasattr(problem, 'feats')
+
+        NW = 1
+        NT = 1
+        NP = 1
+        isbatch = problem.isbatch
+        for k, v in problem.__dict__.items():
+            if k.startswith("worker_"):
+                NW = len(v[0]) if isbatch else len(v)
+            elif k.startswith("task_"):
+                NP = len(v) if isbatch else 1
+                NT = len(v[0]) if isbatch else len(v)
+        NWW = NW * 2
+
+        new_problem = Problem(True)
+        new_problem.feats = {}
+        new_problem.device = 'cpu'
+
+        new_problem.batch_size = NP
+        new_problem.worker_num = NW
+        new_problem.task_num = NT
+
+        new_problem.features = problem.features
+
+        if type(self) == Solver:
+            new_problem.variables = problem.variables
+            new_problem.constraint = problem.constraint
+            new_problem.objective = problem.objective
+            new_problem.environment = problem.environment
+        else:
+            new_problem.variables = []
+            new_problem.constraints = problem.constraints
+            new_problem.oa_estimate_tasks = problem.oa_estimate_tasks
+            new_problem.oa_multiple_steps = problem.oa_multiple_steps
+
+        edge_size_list = ((NWW + NT, NWW + NT), (NW + NT, NW + NT))
+
+        def check_size(f, k, v):
+            assert f, "size error, feature: {}, size: {}".format(k, tuple(v.size()))
+
+        for k, v in problem.__dict__.items():
+            if k == 'solution' and v is not None:
+                v = to_tensor(k, v, isbatch)
+                check_size(v.dim() == 3 and v.size(-1) == 3, k, v)
+            elif k.startswith("worker_task_"):
+                v = to_tensor(k, v, isbatch)
+                check_size(v.dim() in (3, 4) and v.size()[1:3] == (NW, NT), k, v)
+            elif k.startswith("worker_"):
+                v = to_tensor(k, v, isbatch)
+                check_size(v.dim() in (2, 3) and v.size(1) == NW, k, v)
+            elif k.startswith("task_"):
+                v = to_tensor(k, v, isbatch)
+                check_size(v.dim() in (2, 3) and v.size(1) == NT, k, v)
+            elif k.endswith("_matrix"):
+                v = to_tensor(k, v, isbatch)
+                check_size(v.dim() in (3, 4) and v.size()[1:3] in edge_size_list, k, v)
+                if v.size()[1:3] == (NW + NT, NW + NT):
+                    worker_index = torch.arange(NW)
+                    task_index = torch.arange(NT) + NW
+                    index = torch.cat([worker_index, worker_index, task_index])
+                    index1 = index[:, None]
+                    index2 = index[None, :]
+                    v = v[:, index1, index2]
+            elif isinstance(v, np.ndarray):
+                v = torch.tensor(v)
+
+            if isinstance(v, torch.Tensor):
+                new_problem.feats[k] = v
+
+        if pin_memory and self.device.type == 'cuda':
+            new_problem.pin_memory()
+
+        return new_problem
+
+    def to_device(self, problem):
+
+        assert hasattr(problem, 'feats')
+
+        new_problem = copy.copy(problem)
+        new_problem.device = self.device
+        new_problem.feats = {}
+
+        non_blocking = self.device.type == 'cuda'
+        for k, v in problem.feats.items():
+            v = v.to(self.device, non_blocking=non_blocking)
+            new_problem.feats[k] = v
+
+        return new_problem
+
+    def validate(self, problem_list, batch_size):
+        self.agent.eval()
+        with torch.no_grad():
+            valid_result = self.do_validate(problem_list, batch_size)
+
+        self.agent.train()
+        return valid_result
+
+    def do_validate(self, problem_list, batch_size):
+        total_cost1 = 0
+        total_cost2 = 0
+        total_feasible = 0
+        total_problem = 0
+        start_time = time.time()
+        for problem in problem_list:
+            problem = self.to_device(problem)
+            env, _, = self.agent(problem, batch_size)
+
+            NP = problem.batch_size
+            cost = env.cost().sum(1).view(NP, -1)
+            cost1, _ = cost.min(1)
+            cost2 = cost.mean(1)
+            feasible = env.feasible().view(NP, -1)
+            feasible = torch.any(feasible, 1)
+
+            total_cost1 += cost1.sum().item()
+            total_cost2 += cost2.sum().item()
+            total_feasible += feasible.int().sum().item()
+            total_problem += NP
+
+        if dist.is_initialized():
+            data = [total_cost1, total_cost2, total_feasible, total_problem]
+            data = torch.tensor(data, device=self.device)
+            dist.all_reduce(data, op=dist.ReduceOp.SUM)
+            total_cost1, total_cost2, total_feasible, total_problem = data.tolist()
+
+        avg_cost1 = total_cost1 / total_problem
+        avg_cost2 = total_cost2 / total_problem
+        avg_feasible = total_feasible / total_problem
+
+        return avg_cost1, avg_cost2, avg_feasible
+
+    def wrap_dataset(self, dataset, workers, buffers, seed):
+        if isinstance(dataset, IterableDataset):
+            dataset = WrapIterableDataset(dataset, self)
+            dataset = DataLoader(dataset, batch_size=None, pin_memory=True,
+                                 num_workers=workers, prefetch_factor=buffers,
+                                 worker_init_fn=lambda worker_id: torch.manual_seed(seed + worker_id))
+        else:
+            if self.device.type == 'cuda':
+                with torch.cuda.device(cuda_or_none(self.device)):
+                    dataset = WrapDataset(dataset, self)
+                    dataset = DataLoader(dataset, batch_size=None, pin_memory=True, shuffle=True)
+            else:
+                dataset = WrapDataset(dataset, self)
+                dataset = DataLoader(dataset, batch_size=None, pin_memory=True, shuffle=True)
+
+        return dataset
+
+    def print_nn_args(self):
+        for key, value in self.nn_args.items():
+            if type(value) in [int, float, str, bool]:
+                print("nn_args: {} = {}".format(key, value))
+        sys.stdout.flush()
+
+
+def to_tensor(key, value, isbatch):
+    if isinstance(value, torch.Tensor):
+        tensor = value.to('cpu')
+    else:
+        tensor = torch.tensor(value, device='cpu')
+
+    if not isbatch:
+        tensor = tensor[None]
+
+    return tensor
+
+
+def cuda_or_none(device):
+    return device if device.type == 'cuda' else None

greedrl/utils.py

@@ -0,0 +1,65 @@
+import os
+import time
+import torch
+
+act_dict = {}
+act_dict['none'] = lambda x: x
+act_dict['relu'] = torch.relu
+act_dict['tanh'] = torch.tanh
+act_dict['sigmoid'] = torch.sigmoid
+
+
+def get_act(act):
+    return act_dict[act]
+
+
+def to_list(var):
+    if isinstance(var, dict):
+        return {k: to_list(v) for k, v in var.items()}
+    elif isinstance(var, list):
+        return [to_list(v) for v in var]
+    elif isinstance(var, tuple):
+        return (to_list(v) for v in var)
+    elif isinstance(var, torch.Tensor):
+        return var.tolist()
+    else:
+        return var
+
+
+def repeat(tensor, size, dim=0):
+    return tensor.repeat_interleave(size, dim)
+
+
+def get_default_device():
+    if not torch.cuda.is_available():
+        return torch.device("cpu")
+
+    cmd = 'nvidia-smi -q -d Memory | grep -A4 GPU | grep Free'
+    with os.popen(cmd) as result:
+        max_free_mem = 0
+        max_cuda_index = -1
+        for i, line in enumerate(result):
+            free_mem = int(line.strip().split()[2])
+            if free_mem > max_free_mem:
+                max_free_mem = free_mem
+                max_cuda_index = i
+
+    return torch.device("cuda:{}".format(max_cuda_index))
+
+
+def cumem_stats(device, msg):
+    torch.cuda.empty_cache()
+    print("{}, device:{}, memory_allocated: {:.3f}G".format(msg, device,
+                                                            torch.cuda.memory_allocated(device) / (1024 * 1024 * 1024)))
+
+
+cutime_stats_time = None
+
+
+def cutime_stats(device, msg=''):
+    global cutime_stats_time
+    torch.cuda.synchronize(device)
+    if cutime_stats_time is not None:
+        print("{} time: {:.6f}s".format(msg, time.time() - cutime_stats_time))
+
+    cutime_stats_time = time.time()

greedrl/variable.py

@@ -0,0 +1,478 @@
+import torch
+import functools
+
+from .utils import repeat
+
+
+class VarMeta(object):
+    def __init__(self, clazz, **kwargs):
+        self.clazz = clazz
+        self._kwargs = kwargs
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def __call__(self, problem, batch_size, sample_num):
+        kwargs = self._kwargs.copy()
+        kwargs['problem'] = problem.feats
+        kwargs['batch_size'] = batch_size
+        kwargs['sample_num'] = sample_num
+        kwargs['worker_num'] = problem.worker_num
+        kwargs['task_num'] = problem.task_num
+        return self.clazz(**kwargs)
+
+
+def attribute_variable(name, attribute=None):
+    return VarMeta(AttributeVariable, name=name, attribute=attribute)
+
+
+class AttributeVariable:
+    def __init__(self, name, attribute, problem, batch_size, sample_num, worker_num, task_num):
+        if attribute is None:
+            attribute = name;
+
+        self.name = name
+        self.value = problem[attribute]
+
+
+def feature_variable(name, feature=None):
+    return VarMeta(FeatureVariable, name=name, feature=feature)
+
+
+class FeatureVariable:
+    def __init__(self, name, feature, problem, batch_size, sample_num, worker_num, task_num):
+        if feature is None:
+            feature = name
+
+        assert feature == 'id' or feature.startswith("worker_") or feature.startswith("task_")
+
+        self.name = name
+        self.feature = problem[feature]
+        self.value = repeat(self.feature, sample_num)
+
+
+def task_variable(name, feature=None):
+    return VarMeta(TaskVariable, name=name, feature=feature)
+
+
+class TaskVariable:
+    def __init__(self, name, feature, problem, batch_size, sample_num, worker_num, task_num):
+        if feature is None:
+            feature = name
+
+        assert feature.startswith("task_")
+
+        self.name = name
+        self.feature = problem[feature]
+
+        size = list(self.feature.size())
+        size[0] = batch_size
+        del size[1]
+        self.value = self.feature.new_zeros(size)
+
+    def step_task(self, b_index, p_index, t_index):
+        self.value[b_index] = self.feature[p_index, t_index]
+
+
+def worker_variable(name, feature=None):
+    return VarMeta(WorkerVariable, name=name, feature=feature)
+
+
+class WorkerVariable:
+    def __init__(self, name, feature, problem, batch_size, sample_num, worker_num, task_num):
+        if feature is None:
+            feature = name
+
+        assert feature.startswith("worker_")
+
+        self.name = name
+        self.feature = problem[feature]
+
+        size = list(self.feature.size())
+        size[0] = batch_size
+        del size[1]
+        self.value = self.feature.new_zeros(size)
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        self.value[b_index] = self.feature[p_index, w_index]
+
+
+def worker_task_variable(name, feature=None):
+    return VarMeta(WorkerTaskVariable, name=name, feature=feature)
+
+
+class WorkerTaskVariable:
+    def __init__(self, name, feature, problem, batch_size, sample_num, worker_num, task_num):
+        if feature is None:
+            feature = name
+
+        assert feature.startswith("worker_task_")
+
+        self.name = name
+        self.feature = problem[feature]
+
+        size = list(self.feature.size())
+        size[0] = batch_size
+
+        del size[1]
+        self._feature = self.feature.new_zeros(size)
+
+        del size[2]
+        self.value = self.feature.new_zeros(size)
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        self._feature[b_index] = self.feature[p_index, w_index]
+
+    def step_task(self, b_index, p_index, t_index):
+        self.value[b_index] = self._feature[b_index, t_index]
+
+
+def worker_task_group(name, feature=None):
+    return VarMeta(WorkerTaskGroup, name=name, feature=feature)
+
+
+class WorkerTaskGroup:
+    def __init__(self, name, feature, problem, batch_size, sample_num, worker_num, task_num):
+        if feature is None:
+            feature = name
+
+        assert feature.startswith("task_")
+
+        self.name = name
+        self.feature = problem[feature].long()
+
+        NG = self.feature.max() + 1
+        assert torch.all(self.feature >= 0)
+
+        self.value = self.feature.new_zeros(batch_size, NG)
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        self.value[b_index] = 0
+
+    def step_task(self, b_index, p_index, t_index):
+        group = self.feature[p_index, t_index]
+        self.value[b_index, group] += 1;
+
+
+def worker_task_item(name, item_id, item_num):
+    return VarMeta(WorkerTaskItem, name=name, item_id=item_id, item_num=item_num)
+
+
+class WorkerTaskItem:
+    def __init__(self, name, item_id, item_num, problem, batch_size, sample_num, worker_num, task_num):
+        assert item_id.startswith('task_')
+        assert item_num.startswith('task_')
+
+        self.name = name
+        self.item_id = repeat(problem[item_id], sample_num).long()
+        self.item_num = repeat(problem[item_num], sample_num)
+
+        assert torch.all(self.item_id >= 0)
+
+        size = [0, 0]
+        size[0] = self.item_id.size(0)
+        size[1] = self.item_id.max() + 1
+        self.value = self.item_num.new_zeros(size)
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        self.value[b_index] = 0
+
+    def step_task(self, b_index, p_index, t_index):
+        item_id = self.item_id[b_index, t_index]
+        item_num = self.item_num[b_index, t_index]
+        self.value[b_index[:, None], item_id] += item_num
+
+    def make_feat(self):
+        NT = self.item_id.size(1)
+        v = self.value[:, None, :]
+        v = v.expand(-1, NT, -1)
+
+        v = v.gather(2, self.item_id).clamp(0, 1)
+        v = self.item_num.clamp(0, 1) - v
+        return v.clamp(0, 1).sum(2)
+
+
+def task_demand_now(name, feature=None, only_this=False):
+    return VarMeta(TaskDemandNow, name=name, feature=feature, only_this=only_this)
+
+
+class TaskDemandNow:
+    def __init__(self, name, feature, only_this, problem, batch_size, sample_num, worker_num, task_num):
+
+        if feature is None:
+            feature = name
+
+        assert feature.startswith("task_")
+
+        self.name = name
+        self.only_this = only_this
+        self._value = repeat(problem[feature], sample_num)
+
+        assert self._value.dtype in \
+               (torch.int8, torch.int16, torch.int32, torch.int64)
+        assert torch.all(self._value >= 0)
+
+        if only_this:
+            size = self._value.size(0)
+            self.value = self._value.new_zeros(size)
+        else:
+            self.value = self._value
+
+    def step_task(self, b_index, p_index, t_index, done):
+        if done is not None:
+            self._value[b_index, t_index] -= done
+
+        if self.only_this:
+            self.value[b_index] = self._value[b_index, t_index]
+        else:
+            self.value = self._value
+
+
+def worker_count_now(name, feature=None):
+    return VarMeta(WorkerCountNow, name=name, feature=feature)
+
+
+class WorkerCountNow:
+    def __init__(self, name, feature, problem, batch_size, sample_num, worker_num, task_num):
+        if feature is None:
+            feature = name
+
+        assert feature.startswith("worker_")
+
+        self.name = name
+        self.value = repeat(problem[feature], sample_num)
+
+        assert self.value.dtype in \
+               (torch.int8, torch.int16, torch.int32, torch.int64)
+        assert torch.all(self.value >= 0)
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        self.value[b_index, w_index] -= 1
+
+
+def edge_variable(name, feature, last_to_this=False,
+                  this_to_task=False, task_to_end=False, last_to_loop=False):
+    return VarMeta(EdgeVariable, name=name, feature=feature,
+                   last_to_this=last_to_this, this_to_task=this_to_task, task_to_end=task_to_end,
+                   last_to_loop=last_to_loop)
+
+
+class EdgeVariable:
+    def __init__(self, name, feature, last_to_this, this_to_task, task_to_end, last_to_loop,
+                 problem, batch_size, sample_num, worker_num, task_num):
+
+        assert feature.endswith("_matrix")
+
+        flags = [last_to_this, this_to_task, task_to_end, last_to_loop]
+        assert flags.count(True) == 1 and flags.count(False) == 3
+
+        if feature is None:
+            feature = name
+
+        self.name = name
+        self.last_to_this = last_to_this
+        self.this_to_task = this_to_task
+        self.task_to_end = task_to_end
+        self.last_to_loop = last_to_loop
+
+        self.worker_num = worker_num
+        self.task_num = task_num
+
+        self.feature = problem[feature]
+
+        size = list(self.feature.size())
+        size[0] = batch_size
+        del size[1:3]
+
+        if self.this_to_task or self.task_to_end:
+            size.insert(1, task_num)
+            self.value = self.feature.new_zeros(size)
+        else:
+            self.value = self.feature.new_zeros(size)
+
+        self.end_index = self.feature.new_zeros(size[0], dtype=torch.int64)
+        self.loop_index = self.feature.new_zeros(size[0], dtype=torch.int64)
+        self.last_index = self.feature.new_zeros(size[0], dtype=torch.int64)
+        self.task_index = (torch.arange(task_num) + worker_num * 2)[None, :]
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        if self.last_to_this:
+            self.value[b_index] = 0
+            self.last_index[b_index] = w_index
+        elif self.this_to_task:
+            self.do_this_to_task(b_index, p_index, w_index)
+        elif self.task_to_end:
+            self.end_index[b_index] = w_index + self.worker_num
+            self.do_task_to_end(b_index, p_index)
+        elif self.last_to_loop:
+            self.value[b_index] = 0
+            self.last_index[b_index] = w_index
+
+    def step_worker_end(self, b_index, p_index, w_index):
+        this_index = w_index + self.worker_num
+        if self.last_to_this:
+            self.do_last_to_this(b_index, p_index, this_index)
+        elif self.this_to_task:
+            self.do_this_to_task(b_index, p_index, this_index)
+        elif self.task_to_end:
+            pass
+        elif self.last_to_loop:
+            self.do_last_to_loop(b_index, p_index)
+
+    def step_task(self, b_index, p_index, t_index):
+        this_index = t_index + self.worker_num * 2
+        if self.last_to_this:
+            self.do_last_to_this(b_index, p_index, this_index)
+            self.last_index[b_index] = this_index
+        elif self.this_to_task:
+            self.do_this_to_task(b_index, p_index, this_index)
+        elif self.task_to_end:
+            pass
+        elif self.last_to_loop:
+            last_index = self.last_index[b_index]
+            loop_index = self.loop_index[b_index]
+            self.loop_index[b_index] = torch.where(last_index < self.worker_num, this_index, loop_index)
+            self.last_index[b_index] = this_index
+
+    def do_last_to_this(self, b_index, p_index, this_index):
+        last_index = self.last_index[b_index]
+        self.value[b_index] = self.feature[p_index, last_index, this_index]
+
+    def do_this_to_task(self, b_index, p_index, this_index):
+        p_index2 = p_index[:, None]
+        this_index2 = this_index[:, None]
+        task_index2 = self.task_index
+        self.value[b_index] = self.feature[p_index2, this_index2, task_index2]
+
+    def do_task_to_end(self, b_index, p_index):
+        p_index2 = p_index[:, None]
+        task_index2 = self.task_index
+        end_index = self.end_index[b_index]
+        end_index2 = end_index[:, None]
+        self.value[b_index] = self.feature[p_index2, task_index2, end_index2]
+
+    def do_last_to_loop(self, b_index, p_index):
+        loop_index = self.loop_index[b_index]
+        last_index = self.last_index[b_index]
+        self.value[b_index] = self.feature[p_index, last_index, loop_index]
+
+    def make_feat(self):
+        assert self.this_to_task or self.task_to_end, \
+            "one of [this_to_task, task_to_end] must be true"
+        return self.value.clone()
+
+
+def worker_used_resource(name, edge_require=None, task_require=None, task_ready=None, worker_ready=None, task_due=None):
+    return VarMeta(WorkerUsedResource, name=name, edge_require=edge_require, task_require=task_require,
+                   task_ready=task_ready, worker_ready=worker_ready, task_due=task_due)
+
+
+class WorkerUsedResource:
+    def __init__(self, name, edge_require, task_require, task_ready, worker_ready, task_due,
+                 problem, batch_size, sample_num, worker_num, task_num):
+
+        assert edge_require is None or edge_require.endswith("_matrix"), "unsupported edge: {}".format(edge_require)
+        assert task_require is None or task_require.startswith("task_"), "unsupported task_require: {}".format(
+            task_require)
+        assert task_ready is None or task_ready.startswith("task_"), "unsupported task_service: {}".format(task_ready)
+        assert worker_ready is None or worker_ready.startswith("worker_") and not worker_ready.startswith(
+            "worker_task_")
+        assert task_due is None or task_due.startswith("task_"), "unsupported task_due: {}".format(task_due)
+
+        self.name = name
+
+        self.worker_num = worker_num
+        self.task_num = task_num
+
+        if edge_require is None:
+            self.edge_require = None
+        else:
+            self.edge_require = problem[edge_require]
+            self.last_index = self.edge_require.new_zeros(batch_size, dtype=torch.int64)
+
+        if task_require is None:
+            self.task_require = None
+        else:
+            self.task_require = problem[task_require]
+            self.task_require2 = repeat(self.task_require, sample_num)
+
+        if task_ready is None:
+            self.task_ready = None
+        else:
+            self.task_ready = problem[task_ready]
+
+        if worker_ready is None:
+            self.worker_ready = None
+        else:
+            self.worker_ready = problem[worker_ready]
+
+        if task_due is None:
+            self.task_due = None
+        else:
+            self.task_due = problem[task_due]
+
+        tenors = [self.edge_require, self.task_require, self.task_ready, self.worker_ready]
+        tenors = list(filter(lambda x: x is not None, tenors))
+        assert tenors, "at least one of edge_require, task_require, task_ready, worker_ready is required!"
+
+        size = list(tenors[0].size())
+        size[0] = batch_size
+        if self.edge_require is None:
+            del size[1]
+        else:
+            del size[1:3]
+
+        self.value = tenors[0].new_zeros(size)
+
+    def step_worker_start(self, b_index, p_index, w_index):
+        if self.worker_ready is None:
+            self.value[b_index] = 0
+        else:
+            self.value[b_index] = self.worker_ready[p_index, w_index]
+
+        if self.edge_require is not None:
+            self.last_index[b_index] = w_index
+
+    def step_worker_end(self, b_index, p_index, w_index):
+        if self.edge_require is not None:
+            last_index = self.last_index[b_index]
+            this_index = w_index + self.worker_num
+            self.value[b_index] += self.edge_require[p_index, last_index, this_index]
+            self.last_index[b_index] = this_index;
+
+    def step_task(self, b_index, p_index, t_index, done):
+        if done is None:
+            if self.edge_require is not None:
+                last_index = self.last_index[b_index]
+                this_index = t_index + (self.worker_num * 2)
+                self.value[b_index] += self.edge_require[p_index, last_index, this_index]
+                self.last_index[b_index] = this_index
+
+            if self.task_ready is not None:
+                self.value[b_index] = torch.max(self.value[b_index], self.task_ready[p_index, t_index])
+
+        else:
+            if self.task_require is not None:
+                if self.value.dim() == 2:
+                    done = done[:, None]
+                self.value[b_index] += self.task_require[p_index, t_index] * done
+
+    def make_feat(self):
+        assert self.value.dim() == 2, \
+            "value's dim must be 2, actual: {}".format(self.value.dim())
+        assert self.task_require is not None, "task_require is required"
+
+        v = self.value[:, None, :] + self.task_require2
+        return v.clamp(0, 1).sum(2, dtype=v.dtype)
+
+
+def worker_task_sequence(name):
+    return VarMeta(WorkerTaskSequence, name=name)
+
+
+class WorkerTaskSequence:
+    def __init__(self, name, problem, batch_size, sample_num, worker_num, task_num):
+        self.name = name
+        self.value = None
+
+    def step_finish(self, worker_task_seq):
+        self.value = worker_task_seq

requirements.txt

@@ -0,0 +1,7 @@
+torch==1.12.1+cu113
+torchvision==0.13.1+cu113
+torchaudio==0.12.1
+numpy==1.24.2
+Cython==0.29.34
+ortools==9.6.2534
+

setup.py

@@ -0,0 +1,44 @@
+import os
+import re
+import sys
+import time
+import subprocess
+
+from distutils import sysconfig
+from setuptools import setup, Extension, find_packages
+from Cython.Build import build_ext, cythonize
+
+
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=''):
+        Extension.__init__(self, name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext):
+        if isinstance(ext, CMakeExtension):
+            extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
+            if not extdir.endswith(os.path.sep):
+                extdir += os.path.sep
+
+            if not os.path.exists(self.build_temp):
+                os.makedirs(self.build_temp)
+
+            subprocess.check_call(['cmake', ext.sourcedir, '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + extdir],
+                                  cwd=self.build_temp)
+            subprocess.check_call(['cmake', '--build', '.', '--', 'VERBOSE=1', '-j8'], cwd=self.build_temp)
+        else:
+            super().build_extension(ext)
+
+
+ext_modules = [CMakeExtension('greedrl_c')]
+
+setup(
+    name='greedrl',
+    version='1.0.0',
+    packages=find_packages(),
+    ext_modules=ext_modules,
+    cmdclass={'build_ext': CMakeBuild},
+    install_requires=["torch==1.12.1+cu113"],
+)

test/all_test.py

@@ -0,0 +1,7 @@
+from solver_test import *
+from function_test import *
+
+if __name__ == '__main__':
+    
+    unittest.main()
+

test/basetest.py

@@ -0,0 +1,8 @@
+import torch
+import unittest
+
+
+class TestCase(unittest.TestCase):
+    def tearDown(self):
+        torch.cuda.empty_cache()
+

test/function_test.py

@@ -0,0 +1,79 @@
+import sys
+import time
+import torch
+import unittest
+import basetest
+
+from greedrl import Solver
+from greedrl.function import *
+
+device = Solver().device
+
+
+class TestFunction(basetest.TestCase):
+
+    def test_task_group_split(self):
+        group = torch.ones((8, 8), dtype=torch.int32)
+        group[:, 0:4] = 0
+        value = torch.zeros((8, 8), dtype=torch.bool)
+        value[:, 0:4] = True
+        result = task_group_split(group, value)
+        assert not torch.any(result)
+
+        value[:, 0:2] = False
+        result = task_group_split(group, value)
+        assert torch.all(result)
+
+    def test_task_group_split2(self):
+        group = torch.randint(48, (1024, 1000), dtype=torch.int32)
+        value = torch.randint(2, (1024, 1000), dtype=torch.int8) <= 0
+        self.do_test(task_group_split, group, value)
+
+    def test_task_group_priority(self):
+        group = torch.ones((8, 8), dtype=torch.int32)
+        group[:, 0:4] = 0
+        priority = torch.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=torch.int32)
+        priority = priority[None, :].expand(8, -1).clone()
+        value = torch.zeros((8, 8), dtype=torch.bool)
+        value[:, 4:6] = True
+
+        result = task_group_priority(group, priority, value)
+        expected = torch.tensor([False, True, True, True, True, True, False, True])
+        expected = expected[None, :].expand(8, -1)
+        assert torch.all(result == expected)
+
+    def test_task_group_priority2(self):
+        group = torch.randint(48, (1024, 1000), dtype=torch.int32)
+        value = torch.randint(2, (1024, 1000), dtype=torch.int8) < 1
+        priority = torch.randint(2, (1024, 1000), dtype=torch.int32)
+        self.do_test(task_group_priority, group, priority, value)
+
+    def do_test(self, function, *args):
+        print("\ntest {} ...".format(function.__name__))
+        start = time.time()
+        result1 = function(*args)
+        print("time: {:.6f}s, device: {}".format(time.time() - start, args[0].device))
+
+        args = [arg.to(device) for arg in args]
+        result1 = result1.to(device)
+
+        function(*args)
+        self.sync_device(device)
+
+        start = time.time()
+        result2 = function(*args)
+        self.sync_device(device)
+        print("time: {:.6f}s, device: {} ".format(time.time() - start, args[0].device))
+
+        if result1.is_floating_point():
+            assert torch.all(torch.abs(result1 - result2) < 1e-6)
+        else:
+            assert torch.all(result1 == result2)
+
+    def sync_device(self, device):
+        if device.type == 'cuda':
+            torch.cuda.synchronize(device)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/solver_test.py

@@ -0,0 +1,40 @@
+import sys
+import os.path as osp
+import torch
+import unittest
+
+import basetest
+from greedrl import Solver
+from greedrl.const import *
+
+sys.path.append(osp.join(osp.dirname(osp.abspath(__file__)), "../"))
+from examples.cvrp import cvrp
+
+
+class TestSolver(basetest.TestCase):
+    def test(self):
+        problem_list = cvrp.make_problem(1)
+
+        nn_args = {}
+        nn_args['decode_rnn'] = 'GRU'
+        solver = Solver(None, nn_args)
+
+        solver.train(None, problem_list, problem_list,
+                     batch_size=32, max_steps=5, memopt=10)
+
+        solver.train(None, problem_list, problem_list,
+                     batch_size=32, max_steps=5, memopt=10, topk_size=10)
+
+        solver.train(None, problem_list, problem_list,
+                     batch_size=32, max_steps=5, memopt=10, on_policy=False)
+
+        solution = solver.solve(problem_list[0], batch_size=8)
+        assert torch.all(solution.worker_task_sequence[:, -1, 0] == GRL_FINISH)
+        problem_list[0].solution = solution.worker_task_sequence[:, 0:-1, :]
+
+        solution2 = solver.solve(problem_list[0], batch_size=1)
+        assert torch.all(solution.worker_task_sequence == solution2.worker_task_sequence)
+
+
+if __name__ == '__main__':
+    unittest.main()