Unify API for get_state / set_state (#53)

* add api for object

* fix bug

* rename heading_theta/heading

* fix bug

* format

metadrive/base_class/base_object.py

@@ -1,9 +1,16 @@
-from panda3d.bullet import BulletWorld, BulletBodyNode
-from panda3d.core import NodePath
+import math
 
+import numpy as np
 from metadrive.base_class.base_runnable import BaseRunnable
 from metadrive.engine.asset_loader import AssetLoader
 from metadrive.engine.core.physics_world import PhysicsWorld
+from metadrive.utils import Vector
+from metadrive.utils.coordinates_shift import panda_position, metadrive_position, panda_heading, metadrive_heading
+from metadrive.utils.math_utils import clip
+from metadrive.utils.math_utils import norm
+from panda3d.bullet import BulletWorld, BulletBodyNode
+from panda3d.core import LVector3
+from panda3d.core import NodePath
 
 
 class PhysicsNodeList(list):
@@ -129,3 +136,76 @@ def destroy(self):
         self.dynamic_nodes.clear()
         self.static_nodes.clear()
         self._config.clear()
+
+    def set_position(self, position, height=0.4):
+        """
+        Set this object to a place
+        :param position: 2d array or list
+        """
+        self.origin.setPos(panda_position(position, height))
+
+    @property
+    def position(self):
+        return metadrive_position(self.origin.getPos())
+
+    def set_velocity(self, direction: list, value: float):
+        """
+        Set velocity for object including the direction of velocity and the value (speed)
+        The direction of velocity will be normalized automatically, value decided its scale
+        :param position: 2d array or list
+        :param value: speed [m/s]
+        """
+        norm_ratio = value / norm(direction[0], direction[1])
+        self._body.setLinearVelocity(
+            LVector3(direction[0] * norm_ratio, direction[1] * norm_ratio,
+                     self.origin.getPos()[-1])
+        )
+
+    @property
+    def velocity(self):
+        """
+        Velocity, unit: m/s
+        """
+        velocity = self.body.get_linear_velocity()
+        return np.asarray([velocity[0], -velocity[1]])
+
+    @property
+    def speed(self):
+        """
+        return the speed in m/s
+        """
+        velocity = self.body.get_linear_velocity()
+        speed = norm(velocity[0], velocity[1])
+        return clip(speed, 0.0, 100000.0)
+
+    def set_heading_theta(self, heading_theta, rad_to_degree=True) -> None:
+        """
+        Set heading theta for this object
+        :param heading_theta: float
+        :param in_rad: when set to True transfer to degree automatically
+        """
+        h = panda_heading(heading_theta)
+        if rad_to_degree:
+            h = h * 180 / np.pi
+        self.origin.setH(h)
+
+    @property
+    def heading_theta(self):
+        """
+        Get the heading theta of this object, unit [rad]
+        :return:  heading in rad
+        """
+        return metadrive_heading(self.origin.getH()) / 180 * math.pi
+
+    @property
+    def heading(self):
+        """
+        Heading is a vector = [cos(heading_theta), sin(heading_theta)]
+        """
+        real_heading = self.heading_theta
+        # heading = np.array([math.cos(real_heading), math.sin(real_heading)])
+        heading = Vector((math.cos(real_heading), math.sin(real_heading)))
+        return heading
+
+    def set_static(self, flag):
+        self.body.setStatic(flag)

metadrive/component/blocks/tollgate.py

@@ -75,7 +75,7 @@ def _add_building_and_speed_limit(self, road):
                 # add toll
                 position = lane.position(lane.length / 2, 0)
                 building = get_engine().spawn_object(
-                    TollGateBuilding, lane=lane, position=position, heading=lane.heading_at(0)
+                    TollGateBuilding, lane=lane, position=position, heading=lane.heading_theta_at(0)
                 )
                 self.dynamic_nodes.append(building.body)
                 self._block_objects.append(building)

metadrive/component/buildings/tollgate_building.py

@@ -8,14 +8,14 @@ class TollGateBuilding(BaseBuilding):
     BUILDING_LENGTH = 10
     BUILDING_HEIGHT = 5
 
-    def __init__(self, lane, position, heading, random_seed):
-        super(TollGateBuilding, self).__init__(lane, position, heading, random_seed)
+    def __init__(self, lane, position, heading_theta, random_seed):
+        super(TollGateBuilding, self).__init__(lane, position, heading_theta, random_seed)
         air_wall = generate_invisible_static_wall(
             self.BUILDING_LENGTH, lane.width, self.BUILDING_HEIGHT / 2, object_id=self.id
         )
         self.add_body(air_wall)
-        self.origin.setPos(panda_position(position))
-        self.origin.setH(panda_heading(heading))
+        self.set_position(position, 0)
+        self.set_heading_theta(heading)
 
         if self.render:
             building_model = self.loader.loadModel(AssetLoader.file_path("models", "tollgate", "booth.gltf"))

metadrive/component/highway_vehicle/behavior.py

@@ -421,7 +421,7 @@ def steering_features(self, target_lane_index: LaneIndex) -> np.ndarray:
         lane = self.traffic_mgr.current_map.road_network.get_lane(target_lane_index)
         lane_coords = lane.local_coordinates(self.position)
         lane_next_coords = lane_coords[0] + self.speed * self.PURSUIT_TAU
-        lane_future_heading = lane.heading_at(lane_next_coords)
+        lane_future_heading = lane.heading_theta_at(lane_next_coords)
         features = np.array(
             [
                 utils.wrap_to_pi(lane_future_heading - self.heading) * self.LENGTH / utils.not_zero(self.speed),

metadrive/component/highway_vehicle/kinematics.py

@@ -107,7 +107,7 @@ def create_random(
         v = cls(
             traffic_mgr,
             list(lane.position(longitude, 0)),
-            lane.heading_at(longitude),
+            lane.heading_theta_at(longitude),
             speed,
             # random_seed=get_np_random(random_seed)
             # np_random=get_np_random(random_seed)

metadrive/component/lane/abs_lane.py

@@ -47,7 +47,7 @@ def local_coordinates(self, position: np.ndarray) -> Tuple[float, float]:
         raise NotImplementedError()
 
     @abstractmethod
-    def heading_at(self, longitudinal: float) -> float:
+    def heading_theta_at(self, longitudinal: float) -> float:
         """
         Get the lane heading at a given longitudinal lane coordinate.
 
@@ -56,6 +56,10 @@ def heading_at(self, longitudinal: float) -> float:
         """
         raise NotImplementedError()
 
+    def heading_at(self, longitudinal) -> list:
+        heaidng_theta = self.heading_theta_at(longitudinal)
+        return [math.cos(heaidng_theta), math.sin(heaidng_theta)]
+
     @abstractmethod
     def width_at(self, longitudinal: float) -> float:
         """

metadrive/component/lane/circular_lane.py

@@ -48,7 +48,7 @@ def position(self, longitudinal: float, lateral: float) -> Vector:
         # return self.center + (self.radius - lateral * self.direction) * np.array([math.cos(phi), math.sin(phi)])
         return self.center + (self.radius - lateral * self.direction) * Vector((math.cos(phi), math.sin(phi)))
 
-    def heading_at(self, longitudinal: float) -> float:
+    def heading_theta_at(self, longitudinal: float) -> float:
         phi = self.direction * longitudinal / self.radius + self.start_phase
         psi = phi + math.pi / 2 * self.direction
         return psi

metadrive/component/lane/straight_lane.py

@@ -53,7 +53,7 @@ def update_properties(self):
     def position(self, longitudinal: float, lateral: float) -> np.ndarray:
         return self.start + longitudinal * self.direction + lateral * self.direction_lateral
 
-    def heading_at(self, longitudinal: float) -> float:
+    def heading_theta_at(self, longitudinal: float) -> float:
         return self.heading
 
     def width_at(self, longitudinal: float) -> float:

metadrive/component/lane/waypoint_lane.py

@@ -68,7 +68,7 @@ def points_heading(start_p, end_p):
     def width_at(self, longitudinal: float) -> float:
         return self.width
 
-    def heading_at(self, longitudinal: float) -> float:
+    def heading_theta_at(self, longitudinal: float) -> float:
         accumulate_len = 0
         for seg in self.segment_property:
             accumulate_len += seg["length"]

metadrive/component/map/argoverse_map.py

@@ -69,7 +69,7 @@ def _generate(self):
             else:
                 logger.error("Unknown XML item encountered.")
                 raise ValueError("Unknown XML item encountered.")
-        lane_ids = ArgoverseMap.AGMap.get_lane_ids_in_xy_bbox(
+        lane_ids = self.AGMap.get_lane_ids_in_xy_bbox(
             *self.argoverse_position(self.config["center"]), self.config["city"], self.config["radius"]
         )
         self.lane_id_lane = lane_objs

metadrive/component/road/road_network.py

@@ -354,7 +354,8 @@ def position_heading_along_route(self, route: Route, longitudinal: float, latera
         while len(route) > 1 and longitudinal > self.get_lane(route[0]).length:
             longitudinal -= self.get_lane(route[0]).length
             route = route[1:]
-        return self.get_lane(route[0]).position(longitudinal, lateral), self.get_lane(route[0]).heading_at(longitudinal)
+        return self.get_lane(route[0]).position(longitudinal,
+                                                lateral), self.get_lane(route[0]).heading_theta_at(longitudinal)
 
     def get_roads(self, *, direction="all", lane_num=None) -> List:
         """

metadrive/component/static_object/base_static_object.py

@@ -10,24 +10,18 @@
 class BaseStaticObject(BaseObject):
     MASS = 1
 
-    def __init__(self, lane, position: Sequence[float], heading: float = 0., random_seed=None):
+    def __init__(self, lane, position: Sequence[float], heading_theta: float = 0., random_seed=None):
         """
         :param lane: the lane to spawn object
         :param position: cartesian position of object in the surface
-        :param heading: the angle from positive direction of horizontal axis
+        :param heading_theta: the angle from positive direction of horizontal axis
         """
         super(BaseStaticObject, self).__init__(random_seed=random_seed)
-        self.position = position
-        self.speed = 0
-        self.velocity = np.array([0.0, 0.0])
-        self.heading = heading / np.pi * 180
+        self.set_position(position, 0)
+        self.set_heading_theta(heading_theta)
         self.lane_index = lane.index
         self.lane = lane
 
-    @property
-    def heading_theta(self):
-        return self.heading
-
     def set_static(self, static: bool = False):
         mass = 0 if static else self.MASS
         self._body.setMass(mass)

metadrive/component/static_object/traffic_object.py

@@ -28,9 +28,9 @@ def __init__(self, lane, longitude: float, lateral: float, random_seed):
         :param lateral: use to calculate the angle from positive direction of horizontal axis
         """
         position = lane.position(longitude, lateral)
-        heading = lane.heading_at(longitude)
+        heading_theta = lane.heading_theta_at(longitude)
         assert self.NAME is not None, "Assign a name for this class for finding it easily"
-        super(TrafficObject, self).__init__(lane, position, heading, random_seed)
+        super(TrafficObject, self).__init__(lane, position, heading_theta, random_seed)
         self.crashed = False
 
 
@@ -46,8 +46,6 @@ def __init__(self, lane, longitude: float, lateral: float, static: bool = False,
         self.add_body(BaseRigidBodyNode(self.name, self.NAME))
         self.body.addShape(BulletCylinderShape(self.RADIUS, self.HEIGHT))
         self.body.setIntoCollideMask(self.COLLISION_GROUP)
-        self.origin.setPos(panda_position(self.position, self.HEIGHT / 2))
-        self.origin.setH(panda_heading(self.heading))
         if self.render:
             model = self.loader.loadModel(AssetLoader.file_path("models", "traffic_cone", "scene.gltf"))
             model.setScale(0.02)
@@ -68,8 +66,6 @@ def __init__(self, lane, longitude: float, lateral: float, static: bool = False,
         self.add_body(BaseRigidBodyNode(self.name, self.NAME))
         self.body.addShape(BulletCylinderShape(self.RADIUS, self.HEIGHT))
         self.body.setIntoCollideMask(self.COLLISION_GROUP)
-        self.origin.setPos(panda_position(self.position, self.HEIGHT / 2))
-        self.origin.setH(panda_heading(self.heading))
         if self.render:
             model = self.loader.loadModel(AssetLoader.file_path("models", "warning", "warning.gltf"))
             model.setScale(0.02)
@@ -92,8 +88,6 @@ def __init__(self, lane, longitude: float, lateral: float, static: bool = False,
         self.add_body(BaseRigidBodyNode(self.name, self.NAME))
         self.body.addShape(BulletBoxShape((self.WIDTH / 2, self.LENGTH / 2, self.HEIGHT / 2)))
         self.body.setIntoCollideMask(self.COLLISION_GROUP)
-        self.origin.setPos(panda_position(self.position, self.HEIGHT / 2))
-        self.origin.setH(panda_heading(self.heading))
         if self.render:
             model = self.loader.loadModel(AssetLoader.file_path("models", "barrier", "scene.gltf"))
             model.setH(-90)

metadrive/component/vehicle/base_vehicle.py

@@ -306,11 +306,11 @@ def reset(self, random_seed=None, vehicle_config=None, pos: np.ndarray = None, h
         if pos is None:
             lane = map.road_network.get_lane(self.config["spawn_lane_index"])
             pos = lane.position(self.config["spawn_longitude"], self.config["spawn_lateral"])
-            heading = np.rad2deg(lane.heading_at(self.config["spawn_longitude"]))
+            heading = np.rad2deg(lane.heading_theta_at(self.config["spawn_longitude"]))
             self.spawn_place = pos
         heading = -np.deg2rad(heading) - np.pi / 2
         self.set_static(False)
-        self.origin.setPos(panda_position(Vec3(*pos, self.HEIGHT / 2 + 1)))
+        self.set_position(pos, self.HEIGHT / 2 + 1)
         self.origin.setQuat(LQuaternionf(math.cos(heading / 2), 0, 0, math.sin(heading / 2)))
         self.update_map_info(map)
         self.body.clearForces()
@@ -391,8 +391,12 @@ def _dist_to_route_left_right(self):
         return lateral_to_left, lateral_to_right
 
     @property
-    def position(self):
-        return metadrive_position(self.origin.getPos())
+    def heading_theta(self):
+        """
+        Get the heading theta of vehicle, unit [rad]
+        :return:  heading in rad
+        """
+        return (metadrive_heading(self.origin.getH()) - 90) / 180 * math.pi
 
     @property
     def speed(self):
@@ -403,21 +407,6 @@ def speed(self):
         speed = norm(velocity[0], velocity[1]) * 3.6
         return clip(speed, 0.0, 100000.0)
 
-    @property
-    def heading(self):
-        real_heading = self.heading_theta
-        # heading = np.array([math.cos(real_heading), math.sin(real_heading)])
-        heading = Vector((math.cos(real_heading), math.sin(real_heading)))
-        return heading
-
-    @property
-    def heading_theta(self):
-        """
-        Get the heading theta of vehicle, unit [rad]
-        :return:  heading in rad
-        """
-        return (metadrive_heading(self.origin.getH()) - 90) / 180 * math.pi
-
     @property
     def velocity(self) -> np.ndarray:
         return self.speed * self.velocity_direction
@@ -667,21 +656,15 @@ def destroy(self):
         self.image_sensors = None
         self.engine = None
 
-    def set_position(self, position, height=0.4):
+    def set_heading_theta(self, heading_theta, rad_to_degree=True) -> None:
         """
-        Should only be called when restore traffic from episode data
-        :param position: 2d array or list
-        :return: None
-        """
-        self.origin.setPos(panda_position(position, height))
-
-    def set_heading(self, heading_theta) -> None:
-        """
-        Should only be called when restore traffic from episode data
+        Set heading theta for this object
         :param heading_theta: float in rad
-        :return: None
         """
-        self.origin.setH((panda_heading(heading_theta) * 180 / np.pi) - 90)
+        h = panda_heading(heading_theta)
+        if rad_to_degree:
+            h *= 180 / np.pi
+        self.origin.setH(h - 90)
 
     def get_state(self):
         final_road = self.navigation.final_road
@@ -695,7 +678,7 @@ def get_state(self):
         }
 
     def set_state(self, state: dict):
-        self.set_heading(state["heading"])
+        self.set_heading_theta(state["heading"])
         self.set_position(state["position"])
         self._replay_done = state["done"]
         self.set_position(state["position"], height=0.28)
@@ -747,9 +730,6 @@ def arrive_destination(self):
         )
         return flag
 
-    def set_static(self, flag):
-        self.body.setStatic(flag)
-
     @property
     def reference_lanes(self):
         return self.navigation.current_ref_lanes

metadrive/component/vehicle_module/navigation.py

@@ -229,9 +229,9 @@ def _get_info_for_checkpoint(self, lanes_id, lanes, ego_vehicle):
         navi_information.append(clip((ckpt_in_rhs / self.NAVI_POINT_DIST + 1) / 2, 0.0, 1.0))
 
         if lanes_id == 0:
-            lanes_heading = ref_lane.heading_at(ref_lane.local_coordinates(ego_vehicle.position)[0])
+            lanes_heading = ref_lane.heading_theta_at(ref_lane.local_coordinates(ego_vehicle.position)[0])
         else:
-            lanes_heading = ref_lane.heading_at(min(self.PRE_NOTIFY_DIST, ref_lane.length))
+            lanes_heading = ref_lane.heading_theta_at(min(self.PRE_NOTIFY_DIST, ref_lane.length))
 
         # Try to include the current lane's information into the navigation information
         bendradius = 0.0

metadrive/engine/core/chase_camera.py

@@ -162,7 +162,7 @@ def _heading_of_lane(lane, pos: Tuple) -> float:
         :param pos: Tuple, MetaDrive coordinates
         :return: heading theta
         """
-        heading_theta = panda_heading(lane.heading_at(lane.local_coordinates(pos)[0]))
+        heading_theta = panda_heading(lane.heading_theta_at(lane.local_coordinates(pos)[0]))
         return heading_theta
 
     @staticmethod