examples/batching/batching.py

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)