pycityagent 2.0.0a43__cp39-cp39-macosx_11_0_arm64.whl

pycityagent/environment/sim/person_service.py

@@ -0,0 +1,295 @@
+import warnings
+from collections.abc import Awaitable, Coroutine
+from typing import Any, Union, cast
+
+import grpc
+from google.protobuf.json_format import ParseDict
+from mosstool.trip.generator import default_person_template_generator
+from mosstool.util.format_converter import pb2dict
+from pycityproto.city.person.v2 import person_pb2 as person_pb2
+from pycityproto.city.person.v2 import person_service_pb2 as person_service
+from pycityproto.city.person.v2 import person_service_pb2_grpc as person_grpc
+
+from ..utils.protobuf import async_parse
+
+__all__ = ["PersonService"]
+
+
+class PersonService:
+    """
+    交通模拟person服务
+    Traffic simulation person service
+    """
+
+    def __init__(self, aio_channel: grpc.aio.Channel):
+        self._aio_stub = person_grpc.PersonServiceStub(aio_channel)
+
+    @staticmethod
+    def default_person() -> person_pb2.Person:
+        """
+        获取person基本模板
+        Get person basic template
+
+        需要补充的字段有person.home,person.schedules,person.labels
+        The fields that need to be supplemented are person.home, person.schedules, person.labels
+        """
+        person = default_person_template_generator()
+        return person
+
+    @staticmethod
+    def default_dict_person() -> dict:
+        """
+        获取person基本模板，字典格式
+        Get person basic template in dict format.
+
+        需要补充的字段有person.home,person.schedules,person.labels
+        The fields that need to be supplemented are person.home, person.schedules, person.labels
+        """
+        person = default_person_template_generator()
+        return pb2dict(person)
+
+    def GetPerson(
+        self,
+        req: Union[person_service.GetPersonRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], person_service.GetPersonResponse]]:
+        """
+        获取person信息
+        Get person information
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.GetPersonRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.GetPersonResponse
+        """
+        if type(req) != person_service.GetPersonRequest:
+            req = ParseDict(req, person_service.GetPersonRequest())
+        res = cast(
+            Awaitable[person_service.GetPersonResponse], self._aio_stub.GetPerson(req)
+        )
+        return async_parse(res, dict_return)
+
+    def AddPerson(
+        self,
+        req: Union[person_service.AddPersonRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], person_service.AddPersonResponse]]:
+        """
+        新增person
+        Add a new person
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.AddPersonRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.AddPersonResponse
+        """
+        if type(req) != person_service.AddPersonRequest:
+            req = ParseDict(req, person_service.AddPersonRequest())
+        res = cast(
+            Awaitable[person_service.AddPersonResponse], self._aio_stub.AddPerson(req)
+        )
+        return async_parse(res, dict_return)
+
+    def SetSchedule(
+        self,
+        req: Union[person_service.SetScheduleRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], person_service.SetScheduleResponse]]:
+        """
+        修改person的schedule
+        set person's schedule
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.SetScheduleRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.SetScheduleResponse
+        """
+        if type(req) != person_service.SetScheduleRequest:
+            req = ParseDict(req, person_service.SetScheduleRequest())
+        res = cast(
+            Awaitable[person_service.SetScheduleResponse],
+            self._aio_stub.SetSchedule(req),
+        )
+        return async_parse(res, dict_return)
+
+    def GetPersons(
+        self,
+        req: Union[person_service.GetPersonsRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], person_service.GetPersonsResponse]]:
+        """
+        获取多个person信息
+        Get information of multiple persons
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.GetPersonsRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.GetPersonsResponse
+        """
+        if type(req) != person_service.GetPersonsRequest:
+            req = ParseDict(req, person_service.GetPersonsRequest())
+        res = cast(
+            Awaitable[person_service.GetPersonsResponse],
+            self._aio_stub.GetPersons(req),
+        )
+        return async_parse(res, dict_return)
+
+    def GetPersonByLongLatBBox(
+        self,
+        req: Union[person_service.GetPersonByLongLatBBoxRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[
+        Any, Any, Union[dict[str, Any], person_service.GetPersonByLongLatBBoxResponse]
+    ]:
+        """
+        获取特定区域内的person
+        Get persons in a specific region
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.GetPersonByLongLatBBoxRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.GetPersonByLongLatBBoxResponse
+        """
+        if type(req) != person_service.GetPersonByLongLatBBoxRequest:
+            req = ParseDict(req, person_service.GetPersonByLongLatBBoxRequest())
+        res = cast(
+            Awaitable[person_service.GetPersonByLongLatBBoxResponse],
+            self._aio_stub.GetPersonByLongLatBBox(req),
+        )
+        return async_parse(res, dict_return)
+
+    def GetAllVehicles(
+        self,
+        req: Union[person_service.GetAllVehiclesRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[
+        Any, Any, Union[dict[str, Any], person_service.GetAllVehiclesResponse]
+    ]:
+        """
+        获取所有车辆
+        Get all vehicles
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.GetAllVehiclesRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.GetAllVehiclesResponse
+        """
+        if type(req) != person_service.GetAllVehiclesRequest:
+            req = ParseDict(req, person_service.GetAllVehiclesRequest())
+        res = cast(
+            Awaitable[person_service.GetAllVehiclesResponse],
+            self._aio_stub.GetAllVehicles(req),
+        )
+        return async_parse(res, dict_return)
+
+    def ResetPersonPosition(
+        self,
+        req: Union[person_service.ResetPersonPositionRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[
+        Any, Any, Union[dict[str, Any], person_service.ResetPersonPositionResponse]
+    ]:
+        """
+        重置人的位置（将停止当前正在进行的出行，转为sleep状态）
+        Reset person's position (stop the current trip and switch to sleep status)
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.ResetPersonPositionRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.ResetPersonPositionResponse
+        """
+        if type(req) != person_service.ResetPersonPositionRequest:
+            req = ParseDict(req, person_service.ResetPersonPositionRequest())
+        res = cast(
+            Awaitable[person_service.ResetPersonPositionResponse],
+            self._aio_stub.ResetPersonPosition(req),
+        )
+        return async_parse(res, dict_return)
+
+    # RL接口
+
+    def SetControlledVehicleIDs(
+        self,
+        req: Union[person_service.SetControlledVehicleIDsRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[
+        Any, Any, Union[dict[str, Any], person_service.SetControlledVehicleIDsResponse]
+    ]:
+        """
+        设置由外部控制行为的vehicle
+        Set controlled vehicle ID
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.SetControlledVehicleIDsRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.SetControlledVehicleIDsResponse
+        """
+        if type(req) != person_service.SetControlledVehicleIDsRequest:
+            req = ParseDict(req, person_service.SetControlledVehicleIDsRequest())
+        res = cast(
+            Awaitable[person_service.SetControlledVehicleIDsResponse],
+            self._aio_stub.SetControlledVehicleIDs(req),
+        )
+        return async_parse(res, dict_return)
+
+    def FetchControlledVehicleEnvs(
+        self,
+        req: Union[person_service.FetchControlledVehicleEnvsRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[
+        Any,
+        Any,
+        Union[dict[str, Any], person_service.FetchControlledVehicleEnvsResponse],
+    ]:
+        """
+        获取由外部控制行为的vehicle的环境信息
+        Fetch controlled vehicle environment information
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.FetchControlledVehicleEnvsRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.FetchControlledVehicleEnvsResponse
+        """
+        if type(req) != person_service.FetchControlledVehicleEnvsRequest:
+            req = ParseDict(req, person_service.FetchControlledVehicleEnvsRequest())
+        res = cast(
+            Awaitable[person_service.FetchControlledVehicleEnvsResponse],
+            self._aio_stub.FetchControlledVehicleEnvs(req),
+        )
+        return async_parse(res, dict_return)
+
+    def SetControlledVehicleActions(
+        self,
+        req: Union[person_service.SetControlledVehicleActionsRequest, dict],
+        dict_return: bool = True,
+    ) -> Coroutine[
+        Any,
+        Any,
+        Union[dict[str, Any], person_service.SetControlledVehicleActionsResponse],
+    ]:
+        """
+        设置由外部控制行为的vehicle的行为
+        Set controlled vehicle actions
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.person.v2.SetControlledVehicleActionsRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.person.v2.SetControlledVehicleActionsResponse
+        """
+        if type(req) != person_service.SetControlledVehicleActionsRequest:
+            req = ParseDict(req, person_service.SetControlledVehicleActionsRequest())
+        res = cast(
+            Awaitable[person_service.SetControlledVehicleActionsResponse],
+            self._aio_stub.SetControlledVehicleActions(req),
+        )
+        return async_parse(res, dict_return)

pycityagent/environment/sim/road_service.py

@@ -0,0 +1,39 @@
+from collections.abc import Awaitable, Coroutine
+from typing import Any, Union, cast
+
+import grpc
+from google.protobuf.json_format import ParseDict
+from pycityproto.city.map.v2 import road_service_pb2 as road_service
+from pycityproto.city.map.v2 import road_service_pb2_grpc as road_grpc
+
+from ..utils.protobuf import async_parse
+
+__all__ = ["RoadService"]
+
+
+class RoadService:
+    """
+    交通模拟road服务
+    Traffic simulation road service
+    """
+
+    def __init__(self, aio_channel: grpc.aio.Channel):
+        self._aio_stub = road_grpc.RoadServiceStub(aio_channel)
+
+    def GetRoad(
+        self, req: Union[road_service.GetRoadRequest, dict], dict_return: bool = True
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], road_service.GetRoadResponse]]:
+        """
+        查询道路信息
+        Query road information
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.map.v2.GetRoadRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.map.v2.GetRoadResponse
+        """
+        if type(req) != road_service.GetRoadRequest:
+            req = ParseDict(req, road_service.GetRoadRequest())
+        res = cast(Awaitable[road_service.GetRoadResponse], self._aio_stub.GetRoad(req))
+        return async_parse(res, dict_return)

pycityagent/environment/sim/sim_env.py

@@ -0,0 +1,145 @@
+import logging
+import os
+import time
+import atexit
+from subprocess import DEVNULL, Popen
+from typing import Optional
+
+from pycitydata.map import Map
+
+from ..utils import encode_to_base64, find_free_port
+
+__all__ = ["ControlSimEnv"]
+
+
+def _generate_yaml_config(map_file: str, start_step: int, total_step: int) -> str:
+    map_file = os.path.abspath(map_file)
+    return f"""
+input:
+  # 地图
+  map:
+    file: "{map_file}"
+
+control:
+  step:
+    # 模拟器起始步
+    start: {start_step}
+    # 模拟总步数，结束步为起始步+总步数
+    total: {total_step}
+    # 每步的时间间隔
+    interval: 1
+  skip_overtime_trip_when_init: true
+  enable_platoon: false
+  enable_indoor: false
+  prefer_fixed_light: true
+  enable_collision_avoidance: false # 计算性能下降10倍，需要保证subloop>=5
+  enable_go_astray: true # 引入串行的路径规划调用，计算性能下降（幅度不确定）
+  lane_change_model: earliest # mobil （主动变道+强制变道，默认值） earliest （总是尽可能早地变道）
+
+output:
+"""
+
+
+class ControlSimEnv:
+    def __init__(
+        self,
+        task_name: str,
+        map_file: str,
+        # person_file,
+        start_step: int,
+        total_step: int,
+        log_dir: str,
+        min_step_time: int = 1000,
+        timeout: int = 5,
+        simuletgo_addr: Optional[str] = None,
+    ):
+        self._task_name = task_name
+        self._map_file = map_file
+        # self._person_file = person_file
+        self._start_step = start_step
+        self._total_step = total_step
+        self._log_dir = log_dir
+        self._min_step_time = min_step_time
+        self._timeout = timeout
+        self.m = Map(pb_path=map_file)
+        """
+        地图数据
+        """
+
+        # 检查二进制文件是否存在
+        # simulet-go: ~/.local/bin/simulet-go
+        self._simuletgo_path = os.path.expanduser("~/.local/bin/simulet-go")
+        if not os.path.exists(os.path.expanduser(self._simuletgo_path)):
+            raise FileNotFoundError("simulet-go not found, please install it first")
+
+        self._simuletgo_config = _generate_yaml_config(map_file, start_step, total_step)
+        self.simuletgo_port = None
+        self._simuletgo_proc = None
+        self._traffic_client = None
+
+        os.makedirs(log_dir, exist_ok=True)
+
+        self.simuletgo_addr = self.reset(simuletgo_addr)
+
+    def reset(
+        self,
+        simuletgo_addr: Optional[str] = None,
+    ):
+        """
+        Args:
+        - simuletgo_addr: str, simulet-go的地址。如果为None，则启动一个新的simulet-go
+        """
+
+        # 三个地址必须同时为None或者同时不为None
+        if simuletgo_addr is None:
+            # 1. 启动simulet-go
+            # ./simulet-go -config-data configbase64 -job test -syncer http://localhost:53001 -listen :51102
+            assert self.simuletgo_port is None
+            assert self._simuletgo_proc is None
+            self.simuletgo_port = find_free_port()
+            config_base64 = encode_to_base64(self._simuletgo_config)
+            self._simuletgo_proc = Popen(
+                [
+                    self._simuletgo_path,
+                    "-config-data",
+                    config_base64,
+                    "-job",
+                    self._task_name,
+                    "-listen",
+                    f":{self.simuletgo_port}",
+                    "-run.min_step_time",
+                    f"{self._min_step_time}",
+                    "-output",
+                    self._log_dir,
+                    "-cache",
+                    "",
+                    "-log.level",
+                    "error",
+                ],
+                # 忽略输出
+                # stdout=DEVNULL,
+            )
+            logging.info(
+                f"start simulet-go at localhost:{self.simuletgo_port}, PID={self._simuletgo_proc.pid}"
+            )
+            simuletgo_addr = f"http://localhost:{self.simuletgo_port}"
+            atexit.register(self.close)
+            time.sleep(1)
+        elif simuletgo_addr is not None:
+            pass
+        else:
+            # raise ValueError(
+            #     "simuletgo_addr, syncer_addr, routing_addr must be all None or all not None"
+            # )
+            pass
+        return simuletgo_addr
+
+    def close(self):
+        if self._simuletgo_proc is not None:
+            self._simuletgo_proc.terminate()
+            simuletgo_code = self._simuletgo_proc.wait()
+            logging.info(f"simulet-go exit with code {simuletgo_code}")
+
+        self.simuletgo_port = None
+        self._simuletgo_proc = None
+        self._traffic_client = None

pycityagent/environment/sim/social_service.py

@@ -0,0 +1,59 @@
+from collections.abc import Awaitable, Coroutine
+from typing import Any, Union, cast
+
+import grpc
+from google.protobuf.json_format import ParseDict
+from pycityproto.city.social.v1 import social_service_pb2 as social_service
+from pycityproto.city.social.v1 import social_service_pb2_grpc as social_grpc
+
+from ..utils.protobuf import async_parse
+
+__all__ = ["SocialService"]
+
+
+class SocialService:
+    """
+    城市模拟社交服务
+    City simulation social service
+    """
+
+    def __init__(self, aio_channel: grpc.aio.Channel):
+        self._aio_stub = social_grpc.SocialServiceStub(aio_channel)
+
+    def Send(
+        self, req: Union[social_service.SendRequest, dict], dict_return: bool = True
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], social_service.SendResponse]]:
+        """
+        发送消息
+        Send message
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.social.v1.SendRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.social.v1.SendResponse
+        """
+        if type(req) != social_service.SendRequest:
+            req = ParseDict(req, social_service.SendRequest())
+        res = cast(Awaitable[social_service.SendResponse], self._aio_stub.Send(req))
+        return async_parse(res, dict_return)
+
+    def Receive(
+        self, req: Union[social_service.ReceiveRequest, dict], dict_return: bool = True
+    ) -> Coroutine[Any, Any, Union[dict[str, Any], social_service.ReceiveResponse]]:
+        """
+        接收消息
+        Receive message
+
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.social.v1.ReceiveRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.social.v1.ReceiveResponse
+        """
+        if type(req) != social_service.ReceiveRequest:
+            req = ParseDict(req, social_service.ReceiveRequest())
+        res = cast(
+            Awaitable[social_service.ReceiveResponse], self._aio_stub.Receive(req)
+        )
+        return async_parse(res, dict_return)