Vehicle coordinates (#85)

* add convert func wait for test

* coordinates shift and test script

metadrive/component/vehicle/base_vehicle.py

@@ -1,6 +1,7 @@
 import math
+
 from metadrive.utils.space import VehicleParameterSpace, ParameterSpace
-from metadrive.utils.math_utils import safe_clip_for_small_array
+from metadrive.utils.math_utils import norm, wrap_to_pi
 from collections import deque
 from typing import Union, Optional
 import gym
@@ -782,3 +783,34 @@ def attach_to_world(self, parent_node_path, physics_world):
     def set_break_down(self, break_down=True):
         self.break_down = break_down
         # self.set_static(True)
+
+    def convert_to_vehicle_coordinates(self, position, ego_heading=None, ego_position=None):
+        """
+        Give a world position, and convert it to vehicle coordinates
+        The vehicle heading is X direction and right side is Y direction
+        """
+        # Projected to the heading of vehicle
+        pos = ego_heading if ego_position is not None else self.position
+        vector = position - pos
+        forward = self.heading if ego_heading is None else ego_position
+
+        norm_velocity = norm(forward[0], forward[1]) + 1e-6
+        project_on_heading = (vector[0] * forward[0] + vector[1] * forward[1]) / norm_velocity
+
+        side_direction = get_vertical_vector(forward)[1]
+        side_norm = norm(side_direction[0], side_direction[1]) + 1e-6
+        project_on_side = (vector[0] * side_direction[0] + vector[1] * side_direction[1]) / side_norm
+        return project_on_heading, project_on_side
+
+    def convert_to_world_coordinates(self, project_on_heading, project_on_side):
+        """
+        Give a position in vehicle coordinates, and convert it to world coordinates
+        The vehicle heading is X direction and right side is Y direction
+        """
+        theta = np.arctan2(project_on_side, project_on_heading)
+        theta = wrap_to_pi(self.heading_theta) + wrap_to_pi(theta)
+        norm_len = norm(project_on_heading, project_on_side)
+        position = self.position
+        heading = np.sin(theta) * norm_len
+        side = np.cos(theta) * norm_len
+        return position[0] + side, position[1] + heading

metadrive/tests/test_functionality/test_vehicle_coordinates.py

@@ -0,0 +1,41 @@
+from metadrive.envs.metadrive_env import MetaDriveEnv
+
+
+def test_coordinates_shift():
+    try:
+        env = MetaDriveEnv(
+            {
+                "environment_num": 100,
+                "traffic_density": .0,
+                "traffic_mode": "trigger",
+                "start_seed": 22,
+                # "manual_control": True,
+                # "use_render": True,
+                "decision_repeat": 5,
+                "rgb_clip": True,
+                "pstats": True,
+                # "discrete_action": True,
+                "map": "SSSSSS",
+            }
+        )
+        env.reset()
+        env.vehicle.set_velocity([1, 0], 10)
+        print(env.vehicle.speed)
+        pos = [(x, y) for x in [-10, 0, 10] for y in [-20, 0, 20]] * 10
+        p = pos.pop()
+        for s in range(1, 100000):
+            o, r, d, info = env.step([1, 0.3])
+            if s % 10 == 0:
+                if len(pos) == 0:
+                    break
+                p = pos.pop()
+            heading, side = env.vehicle.convert_to_vehicle_coordinates(p)
+            recover_pos = env.vehicle.convert_to_world_coordinates(heading, side)
+            if abs(recover_pos[0] - p[0]) + abs(recover_pos[1] - p[1]) > 0.1:
+                raise ValueError("vehicle coordinates convert error!")
+    finally:
+        env.close()
+
+
+if __name__ == "__main__":
+    test_coordinates_shift()

metadrive/tests/vis_env/vis_metadrive_env.py

@@ -14,7 +14,7 @@
             "global_light": True,
             # "debug_static_world":True,
             "cull_scene": False,
-            "offscreen_render": True,
+            # "offscreen_render": True,
             # "controller": "joystick",
             "manual_control": True,
             "use_render": True,
@@ -38,7 +38,7 @@
             "driving_reward": 1.0,
             "vehicle_config": {
                 "enable_reverse": True,
-                "image_source": "depth_camera",
+                # "image_source": "depth_camera",
                 # "random_color": True
                 # "show_lidar": True,
                 # "spawn_lane_index":("1r1_0_", "1r1_1_", 0),