From 8fffccd9c034a149ca71b74010c71e1f2fabf8d0 Mon Sep 17 00:00:00 2001 From: Zhenghao Peng Date: Mon, 2 Dec 2024 18:01:55 -0800 Subject: [PATCH] SimGen-related code (#777) * non-breaking change * non-breaking change * non-breaking changes: Introduce PedestrianBoundingBox * non-breaking changes * non-breaking changes * Fix a bug in cyclist: The object shape is wrong. * non-breaking changes: Introduce CyclistBoundingBox * non-breaking changes: Introduce VaryingDynamicsBoundingBoxVehicle (TODO: Semantics.BUS.label is not used) * non-breaking changes: Introduce drive_in_real_env_with_bounding_box.py (example file) * Fix a bug in ReplayPolicy: The object is now set_static(True) so it won't be affected by gravity or collision. * non-breaking change * non-breaking changes: implement generate_distinct_rgb_values (for instance level colors?) * Introduce config "use_bounding_box" in ScenarioEnv, impl the proper config postprocessing * non-breaking changes: Introduce ScenarioAgentManager * non-breaking changes: Update Instance Camera * non-breaking changes: fix a bug * Update scenario_traffic_manager.py: Use Cyc/PedBoundingBox, deal with use_bounding_box * Update scenario_traffic_manager.py: add w/l/h info into vehicle_config before spawning it. * non-breaking changes: Update scenario_map_manager.py, add w/l/h to vehicle config (traffic agent) * Update get_vehicle_type, remove randomization of vehicle type in ScenarioEnv * Default even_sample_vehicle_class=False. (remove randomization of veh type in ScenarioEnv) * non-breaking change * Create new car model with different scale and shape, if self.config['scale'] is set. * Deprecate even_sample_vehicle_class in ScenarioEnv * Allow to update the tire_scale if config['scale'] is set in ScenarioEnv * Fix a bug in ScenarioMapManager: for Nuscene's data, the width / length is wrong (error from source data, not our problem). * format * Minor update: change API for camera.get_image --- .../trajectory_navigation.py | 17 +- metadrive/component/sensors/base_camera.py | 11 +- .../component/sensors/instance_camera.py | 25 +- .../component/sensors/semantic_camera.py | 2 +- .../base_traffic_participant.py | 4 +- .../component/traffic_participants/cyclist.py | 191 +++++++++++++- .../traffic_participants/pedestrian.py | 177 ++++++++++++- metadrive/component/vehicle/base_vehicle.py | 53 +++- metadrive/component/vehicle/vehicle_type.py | 211 ++++++++++++--- metadrive/constants.py | 10 +- metadrive/engine/base_engine.py | 42 +-- metadrive/engine/core/engine_core.py | 4 +- metadrive/engine/core/main_camera.py | 14 +- metadrive/engine/core/physics_world.py | 6 +- metadrive/envs/base_env.py | 3 + metadrive/envs/scenario_env.py | 23 +- .../drive_in_real_env_with_bounding_box.py | 61 +++++ metadrive/manager/scenario_agent_manager.py | 8 + metadrive/manager/scenario_map_manager.py | 30 ++- metadrive/manager/scenario_traffic_manager.py | 138 ++++++---- metadrive/policy/replay_policy.py | 1 + metadrive/scenario/parse_object_state.py | 8 +- metadrive/scenario/utils.py | 23 +- .../test_component/test_lane_line_detector.py | 246 +++++++++--------- .../test_export_scenario.py | 4 +- .../test_export_scenario_consistency_test.py | 4 +- .../tests/test_sensors/test_instance_cam.py | 14 +- 27 files changed, 1050 insertions(+), 280 deletions(-) create mode 100755 metadrive/examples/drive_in_real_env_with_bounding_box.py create mode 100644 metadrive/manager/scenario_agent_manager.py diff --git a/metadrive/component/navigation_module/trajectory_navigation.py b/metadrive/component/navigation_module/trajectory_navigation.py index 308c159ae..e36482391 100644 --- a/metadrive/component/navigation_module/trajectory_navigation.py +++ b/metadrive/component/navigation_module/trajectory_navigation.py @@ -33,9 +33,9 @@ def __init__( if show_dest_mark or show_line_to_dest: get_logger().warning("show_dest_mark and show_line_to_dest are not supported in TrajectoryNavigation") super(TrajectoryNavigation, self).__init__( - show_navi_mark=False, - show_dest_mark=False, - show_line_to_dest=False, + show_navi_mark=show_navi_mark, + show_dest_mark=show_dest_mark, + show_line_to_dest=show_line_to_dest, panda_color=panda_color, name=name, vehicle_config=vehicle_config @@ -145,6 +145,17 @@ def update_localization(self, ego_vehicle): # Use RC as the only criterion to determine arrival in Scenario env. self._route_completion = long / self.reference_trajectory.length + if self._show_navi_info: + # Whether to visualize little boxes in the scene denoting the checkpoints + pos_of_goal = ckpts[1] + self._goal_node_path.setPos(panda_vector(pos_of_goal[0], pos_of_goal[1], self.MARK_HEIGHT)) + self._goal_node_path.setH(self._goal_node_path.getH() + 3) + # self.navi_arrow_dir = [lanes_heading1, lanes_heading2] + dest_pos = self._dest_node_path.getPos() + self._draw_line_to_dest(start_position=ego_vehicle.position, end_position=(dest_pos[0], dest_pos[1])) + navi_pos = self._goal_node_path.getPos() + self._draw_line_to_navi(start_position=ego_vehicle.position, end_position=(navi_pos[0], navi_pos[1])) + def get_current_lateral_range(self, current_position, engine) -> float: return self.current_lane.width * 2 diff --git a/metadrive/component/sensors/base_camera.py b/metadrive/component/sensors/base_camera.py index 752573926..c0a88f8c2 100644 --- a/metadrive/component/sensors/base_camera.py +++ b/metadrive/component/sensors/base_camera.py @@ -97,7 +97,7 @@ def _make_cuda_texture(self): def enable_cuda(self): return self is not None and self._enable_cuda - def get_image(self, base_object): + def get_image(self, base_object, mode="bgr"): """ Put camera to an object and get the image. """ @@ -107,13 +107,18 @@ def get_image(self, base_object): self.cam.reparentTo(base_object.origin) img = self.get_rgb_array_cpu() self.track(original_parent, original_position, original_hpr) - return img + if mode == "bgr": + return img + elif mode == "rgb": + return img[..., ::-1] + else: + raise ValueError("Unknown mode: {}".format(mode)) def save_image(self, base_object, name="debug.png"): """ Put camera to an object and save the image to the disk """ - img = self.get_image(base_object) + img = self.get_image(base_object, mode="bgr") cv2.imwrite(name, img) def track(self, new_parent_node: NodePath, position, hpr): diff --git a/metadrive/component/sensors/instance_camera.py b/metadrive/component/sensors/instance_camera.py index 9adb6a780..76e0f7296 100644 --- a/metadrive/component/sensors/instance_camera.py +++ b/metadrive/component/sensors/instance_camera.py @@ -1,4 +1,8 @@ -from panda3d.core import RenderState, LightAttrib, ColorAttrib, ShaderAttrib, TextureAttrib, FrameBufferProperties +from typing import Union + +import numpy as np +from panda3d.core import NodePath +from panda3d.core import RenderState, LightAttrib, ColorAttrib, ShaderAttrib, TextureAttrib from metadrive.component.sensors.base_camera import BaseCamera from metadrive.constants import CamMask @@ -21,11 +25,16 @@ def track(self, new_parent_node, position, hpr): self._setup_effect() super().track(new_parent_node, position, hpr) + def perceive( + self, to_float=True, new_parent_node: Union[NodePath, None] = None, position=None, hpr=None + ) -> np.ndarray: + self._setup_effect() + return super().perceive(to_float, new_parent_node, position, hpr) + def _setup_effect(self): """ - Use tag to apply color to different object class + Use tag to apply color to different objects (instances) Returns: None - """ # setup camera @@ -34,6 +43,11 @@ def _setup_effect(self): else: mapping = get_engine().id_c spawned_objects = get_engine().get_objects() + + ##Ensure consistency between color mapping and the objects actually active in the engine. + mapping_set, object_set = set(list(mapping.keys())), set(list(spawned_objects.keys())) + assert (len(mapping_set.difference(object_set)) == 0) + for id, obj in spawned_objects.items(): obj.origin.setTag(CameraTagStateKey.ID, id) cam = self.get_cam().node() @@ -44,5 +58,8 @@ def _setup_effect(self): ColorAttrib.makeFlat((0, 0, 0, 1)), 1 ) ) - for id, c in mapping.items(): + + for id in spawned_objects.keys(): + c = mapping[id] + assert c in self.engine.COLORS_OCCUPIED cam.setTagState(id, RenderState.make(ColorAttrib.makeFlat((c[0], c[1], c[2], 1)), 1)) diff --git a/metadrive/component/sensors/semantic_camera.py b/metadrive/component/sensors/semantic_camera.py index de4a4f0b7..480cbb5c9 100644 --- a/metadrive/component/sensors/semantic_camera.py +++ b/metadrive/component/sensors/semantic_camera.py @@ -38,7 +38,7 @@ def _setup_effect(self): ) else: - if label == Semantics.PEDESTRIAN.label: + if label == Semantics.PEDESTRIAN.label and not self.engine.global_config.get("use_bounding_box", False): # PZH: This is a workaround fix to make pedestrians animated. cam.setTagState( label, diff --git a/metadrive/component/traffic_participants/base_traffic_participant.py b/metadrive/component/traffic_participants/base_traffic_participant.py index 6406ec03f..6f2669e9a 100644 --- a/metadrive/component/traffic_participants/base_traffic_participant.py +++ b/metadrive/component/traffic_participants/base_traffic_participant.py @@ -14,8 +14,8 @@ class BaseTrafficParticipant(BaseObject): COLLISION_MASK = CollisionGroup.TrafficParticipants HEIGHT = None - def __init__(self, position: Sequence[float], heading_theta: float = 0., random_seed=None, name=None): - super(BaseTrafficParticipant, self).__init__(random_seed=random_seed, name=name) + def __init__(self, position: Sequence[float], heading_theta: float = 0., random_seed=None, name=None, config=None): + super(BaseTrafficParticipant, self).__init__(random_seed=random_seed, name=name, config=config) self.set_position(position, self.HEIGHT / 2 if hasattr(self, "HEIGHT") else 0) self.set_heading_theta(heading_theta) assert self.MASS is not None, "No mass for {}".format(self.class_name) diff --git a/metadrive/component/traffic_participants/cyclist.py b/metadrive/component/traffic_participants/cyclist.py index 036271a8e..0d2f5db07 100644 --- a/metadrive/component/traffic_participants/cyclist.py +++ b/metadrive/component/traffic_participants/cyclist.py @@ -1,11 +1,9 @@ -from metadrive.component.traffic_participants.base_traffic_participant import BaseTrafficParticipant -from typing import Tuple - from panda3d.bullet import BulletBoxShape -from panda3d.bullet import BulletCylinderShape +from panda3d.core import LineSegs, NodePath -from metadrive.constants import MetaDriveType, Semantics +from metadrive.component.traffic_participants.base_traffic_participant import BaseTrafficParticipant from metadrive.constants import CollisionGroup +from metadrive.constants import MetaDriveType, Semantics from metadrive.engine.asset_loader import AssetLoader from metadrive.engine.physics_node import BaseRigidBodyNode @@ -19,13 +17,29 @@ class Cyclist(BaseTrafficParticipant): HEIGHT = 1.75 - def __init__(self, position, heading_theta, random_seed, name=None): + DEFAULT_LENGTH = 1.75 # meters + DEFAULT_HEIGHT = 1.75 # meters + DEFAULT_WIDTH = 0.4 # meters + + @property + def LENGTH(self): + return self.DEFAULT_LENGTH + + @property + def HEIGHT(self): + return self.DEFAULT_HEIGHT + + @property + def WIDTH(self): + return self.DEFAULT_WIDTH + + def __init__(self, position, heading_theta, random_seed, name=None, **kwargs): super(Cyclist, self).__init__(position, heading_theta, random_seed, name=name) self.set_metadrive_type(self.TYPE_NAME) n = BaseRigidBodyNode(self.name, self.TYPE_NAME) self.add_body(n) - self.body.addShape(BulletBoxShape((self.LENGTH / 2, self.WIDTH / 2, self.HEIGHT / 2))) + self.body.addShape(BulletBoxShape((self.WIDTH / 2, self.LENGTH / 2, self.HEIGHT / 2))) if self.render: if Cyclist.MODEL is None: model = self.loader.loadModel(AssetLoader.file_path("models", "bicycle", "scene.gltf")) @@ -38,10 +52,169 @@ def set_velocity(self, direction, value=None, in_local_frame=False): super(Cyclist, self).set_velocity(direction, value, in_local_frame) self.standup() + def get_state(self): + state = super(Cyclist, self).get_state() + state.update({ + "length": self.LENGTH, + "width": self.WIDTH, + "height": self.HEIGHT, + }) + return state + + +class CyclistBoundingBox(BaseTrafficParticipant): + MASS = 80 # kg + TYPE_NAME = MetaDriveType.CYCLIST + COLLISION_MASK = CollisionGroup.TrafficParticipants + SEMANTIC_LABEL = Semantics.BIKE.label + + # for random color choosing + MATERIAL_COLOR_COEFF = 1.6 # to resist other factors, since other setting may make color dark + MATERIAL_METAL_COEFF = 0.1 # 0-1 + MATERIAL_ROUGHNESS = 0.8 # smaller to make it more smooth, and reflect more light + MATERIAL_SHININESS = 128 # 0-128 smaller to make it more smooth, and reflect more light + MATERIAL_SPECULAR_COLOR = (3, 3, 3, 3) + + def __init__(self, position, heading_theta, random_seed, name=None, **kwargs): + config = {"width": kwargs["width"], "length": kwargs["length"], "height": kwargs["height"]} + # config = {"width": kwargs["length"], "length": kwargs["width"], "height": kwargs["height"]} + super(CyclistBoundingBox, self).__init__(position, heading_theta, random_seed, name=name, config=config) + self.set_metadrive_type(self.TYPE_NAME) + n = BaseRigidBodyNode(self.name, self.TYPE_NAME) + self.add_body(n) + + self.body.addShape(BulletBoxShape((self.WIDTH / 2, self.LENGTH / 2, self.HEIGHT / 2))) + if self.render: + model = AssetLoader.loader.loadModel(AssetLoader.file_path("models", "box.bam")) + model.setScale((self.WIDTH, self.LENGTH, self.HEIGHT)) + model.setTwoSided(False) + self._instance = model.instanceTo(self.origin) + + # Add some color to help debug + from panda3d.core import Material, LVecBase4 + import seaborn as sns + + show_contour = self.config["show_contour"] if "show_contour" in self.config else False + if show_contour: + # ========== Draw the contour of the bounding box ========== + # Draw the bottom of the car first + line_seg = LineSegs("bounding_box_contour1") + zoffset = model.getZ() + line_seg.setThickness(2) + line_color = [0.0, 0.0, 0.0] + out_offset = 0.02 + w = self.WIDTH / 2 + out_offset + l = self.LENGTH / 2 + out_offset + h = self.HEIGHT / 2 + out_offset + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(-w, l, h + zoffset) + line_seg.drawTo(-w, l, -h + zoffset) + line_seg.drawTo(w, l, -h + zoffset) + line_seg.drawTo(w, l, h + zoffset) + + # draw cross line + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_seg.moveTo(w, -l, h + zoffset) + line_seg.drawTo(w, l, -h + zoffset) + + line_seg.moveTo(w, -l, h + zoffset) + line_seg.drawTo(-w, -l, h + zoffset) + line_seg.drawTo(-w, -l, -h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_seg.drawTo(w, -l, h + zoffset) + + # draw vertical & horizontal line + line_seg.moveTo(-w, l, 0 + zoffset) + line_seg.drawTo(-w, -l, 0 + zoffset) + line_seg.moveTo(-w, 0, h + zoffset) + line_seg.drawTo(-w, 0, -h + zoffset) + + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(w, -l, h + zoffset) + line_seg.moveTo(-w, l, h + zoffset) + line_seg.drawTo(-w, -l, h + zoffset) + line_seg.moveTo(-w, l, -h + zoffset) + line_seg.drawTo(-w, -l, -h + zoffset) + line_seg.moveTo(w, l, -h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_np = NodePath(line_seg.create(True)) + line_material = Material() + line_material.setBaseColor(LVecBase4(*line_color[:3], 1)) + line_np.setMaterial(line_material, True) + line_np.reparentTo(self.origin) + + color = sns.color_palette("colorblind") + color.remove(color[2]) # Remove the green and leave it for special vehicle + idx = 0 + rand_c = color[idx] + rand_c = (1.0, 0.0, 0.0) + self._panda_color = rand_c + material = Material() + material.setBaseColor( + ( + self.panda_color[0] * self.MATERIAL_COLOR_COEFF, self.panda_color[1] * self.MATERIAL_COLOR_COEFF, + self.panda_color[2] * self.MATERIAL_COLOR_COEFF, 0. + ) + ) + material.setMetallic(self.MATERIAL_METAL_COEFF) + material.setSpecular(self.MATERIAL_SPECULAR_COLOR) + material.setRefractiveIndex(1.5) + material.setRoughness(self.MATERIAL_ROUGHNESS) + material.setShininess(self.MATERIAL_SHININESS) + material.setTwoside(False) + self.origin.setMaterial(material, True) + + def reset(self, position, heading_theta: float = 0., random_seed=None, name=None, *args, **kwargs): + super(CyclistBoundingBox, self).reset(position, heading_theta, random_seed, name, *args, **kwargs) + config = {"width": kwargs["width"], "length": kwargs["length"], "height": kwargs["height"]} + self.update_config(config) + if self._instance is not None: + self._instance.detachNode() + if self.render: + model = AssetLoader.loader.loadModel(AssetLoader.file_path("models", "box.bam")) + model.setScale((self.WIDTH, self.LENGTH, self.HEIGHT)) + model.setTwoSided(False) + self._instance = model.instanceTo(self.origin) + + # Add some color to help debug + from panda3d.core import Material, LVecBase4 + import seaborn as sns + color = sns.color_palette("colorblind") + color.remove(color[2]) # Remove the green and leave it for special vehicle + idx = 0 + rand_c = color[idx] + rand_c = (1.0, 0.0, 0.0) + self._panda_color = rand_c + material = Material() + material.setBaseColor( + ( + self.panda_color[0] * self.MATERIAL_COLOR_COEFF, self.panda_color[1] * self.MATERIAL_COLOR_COEFF, + self.panda_color[2] * self.MATERIAL_COLOR_COEFF, 0. + ) + ) + material.setMetallic(self.MATERIAL_METAL_COEFF) + material.setSpecular(self.MATERIAL_SPECULAR_COLOR) + material.setRefractiveIndex(1.5) + material.setRoughness(self.MATERIAL_ROUGHNESS) + material.setShininess(self.MATERIAL_SHININESS) + material.setTwoside(False) + self.origin.setMaterial(material, True) + + def set_velocity(self, direction, value=None, in_local_frame=False): + super(CyclistBoundingBox, self).set_velocity(direction, value, in_local_frame) + self.standup() + @property def WIDTH(self): - return 0.4 + # return self.config["width"] + return self.config["length"] + + @property + def HEIGHT(self): + return self.config["height"] @property def LENGTH(self): - return 1.75 + # return self.config["length"] + return self.config["width"] diff --git a/metadrive/component/traffic_participants/pedestrian.py b/metadrive/component/traffic_participants/pedestrian.py index fed51f295..906ec53f1 100644 --- a/metadrive/component/traffic_participants/pedestrian.py +++ b/metadrive/component/traffic_participants/pedestrian.py @@ -1,6 +1,10 @@ +from typing import Dict + from direct.actor.Actor import Actor -from panda3d.bullet import BulletCylinderShape +from panda3d.bullet import BulletCylinderShape, BulletBoxShape +from panda3d.core import LVecBase4 from panda3d.core import LVector3 +from panda3d.core import LineSegs, NodePath from metadrive.component.traffic_participants.base_traffic_participant import BaseTrafficParticipant from metadrive.constants import MetaDriveType, Semantics @@ -21,7 +25,7 @@ class Pedestrian(BaseTrafficParticipant): # SPEED_LIST = [0.6, 1.2, 2.2] Too much speed choice jeopardise the performance SPEED_LIST = [0.4, 1.2] - def __init__(self, position, heading_theta, random_seed=None, name=None): + def __init__(self, position, heading_theta, random_seed=None, name=None, *args, **kwargs): super(Pedestrian, self).__init__(position, heading_theta, random_seed, name=name) self.set_metadrive_type(self.TYPE_NAME) # self.origin.setDepthOffset(1) @@ -116,3 +120,172 @@ def top_down_width(self): @property def top_down_length(self): return self.RADIUS * 2 + + def get_state(self) -> Dict: + state = super(Pedestrian, self).get_state() + state.update( + { + "length": self.RADIUS * 2, + "width": self.RADIUS * 2, + "height": self.HEIGHT, + "radius": self.RADIUS, + } + ) + return state + + +class PedestrianBoundingBox(BaseTrafficParticipant): + MASS = 70 # kg + TYPE_NAME = MetaDriveType.PEDESTRIAN + SEMANTIC_LABEL = Semantics.PEDESTRIAN.label + + # for random color choosing + MATERIAL_COLOR_COEFF = 1.6 # to resist other factors, since other setting may make color dark + MATERIAL_METAL_COEFF = 0.1 # 0-1 + MATERIAL_ROUGHNESS = 0.8 # smaller to make it more smooth, and reflect more light + MATERIAL_SHININESS = 128 # 0-128 smaller to make it more smooth, and reflect more light + MATERIAL_SPECULAR_COLOR = (3, 3, 3, 3) + + def __init__(self, position, heading_theta, width, length, height, random_seed=None, name=None): + config = {} + config["width"] = width + config["length"] = length + config["height"] = height + + super(PedestrianBoundingBox, self).__init__(position, heading_theta, random_seed, name=name, config=config) + self.set_metadrive_type(self.TYPE_NAME) + n = BaseRigidBodyNode(self.name, self.TYPE_NAME) + self.add_body(n) + + # PZH: Use BoxShape instead of CylinderShape + # self.body.addShape(BulletCylinderShape(self.RADIUS, self.HEIGHT)) + self.body.addShape(BulletBoxShape(LVector3(width / 2, length / 2, height / 2))) + + self._instance = None + if self.render: + # PZH: Load a box model and resize it to the vehicle size + model = AssetLoader.loader.loadModel(AssetLoader.file_path("models", "box.bam")) + model.setScale((self.WIDTH, self.LENGTH, self.HEIGHT)) + model.setTwoSided(False) + self._instance = model.instanceTo(self.origin) + + # Add some color to help debug + from panda3d.core import Material + import seaborn as sns + + show_contour = self.config["show_contour"] if "show_contour" in self.config else False + if show_contour: + # ========== Draw the contour of the bounding box ========== + # Draw the bottom of the car first + line_seg = LineSegs("bounding_box_contour1") + zoffset = model.getZ() + line_seg.setThickness(2) + line_color = [0.0, 0.0, 1.0] + out_offset = 0.02 + w = self.WIDTH / 2 + out_offset + l = self.LENGTH / 2 + out_offset + h = self.HEIGHT / 2 + out_offset + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(-w, l, h + zoffset) + line_seg.drawTo(-w, l, -h + zoffset) + line_seg.drawTo(w, l, -h + zoffset) + line_seg.drawTo(w, l, h + zoffset) + + # draw cross line + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_seg.moveTo(w, -l, h + zoffset) + line_seg.drawTo(w, l, -h + zoffset) + + line_seg.moveTo(w, -l, h + zoffset) + line_seg.drawTo(-w, -l, h + zoffset) + line_seg.drawTo(-w, -l, -h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_seg.drawTo(w, -l, h + zoffset) + + # draw vertical & horizontal line + line_seg.moveTo(-w, l, 0 + zoffset) + line_seg.drawTo(-w, -l, 0 + zoffset) + line_seg.moveTo(-w, 0, h + zoffset) + line_seg.drawTo(-w, 0, -h + zoffset) + + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(w, -l, h + zoffset) + line_seg.moveTo(-w, l, h + zoffset) + line_seg.drawTo(-w, -l, h + zoffset) + line_seg.moveTo(-w, l, -h + zoffset) + line_seg.drawTo(-w, -l, -h + zoffset) + line_seg.moveTo(w, l, -h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_np = NodePath(line_seg.create(True)) + line_material = Material() + line_material.setBaseColor(LVecBase4(*line_color[:3], 1)) + line_np.setMaterial(line_material, True) + line_np.reparentTo(self.origin) + + color = sns.color_palette("colorblind") + color.remove(color[2]) # Remove the green and leave it for special vehicle + idx = 0 + rand_c = color[idx] + rand_c = (0.0, 1.0, 0.0) + self._panda_color = rand_c + material = Material() + material.setBaseColor( + ( + self.panda_color[0] * self.MATERIAL_COLOR_COEFF, self.panda_color[1] * self.MATERIAL_COLOR_COEFF, + self.panda_color[2] * self.MATERIAL_COLOR_COEFF, 0. + ) + ) + material.setMetallic(self.MATERIAL_METAL_COEFF) + material.setSpecular(self.MATERIAL_SPECULAR_COLOR) + material.setRefractiveIndex(1.5) + material.setRoughness(self.MATERIAL_ROUGHNESS) + material.setShininess(self.MATERIAL_SHININESS) + material.setTwoside(False) + self.origin.setMaterial(material, True) + + def reset(self, position, heading_theta: float = 0., random_seed=None, name=None, *args, **kwargs): + super(PedestrianBoundingBox, self).reset(position, heading_theta, random_seed, name, *args, **kwargs) + config = {"width": kwargs["width"], "length": kwargs["length"], "height": kwargs["height"]} + self.update_config(config) + if self._instance is not None: + self._instance.detachNode() + if self.render: + model = AssetLoader.loader.loadModel(AssetLoader.file_path("models", "box.bam")) + model.setScale((self.WIDTH, self.LENGTH, self.HEIGHT)) + model.setTwoSided(False) + self._instance = model.instanceTo(self.origin) + + def set_velocity(self, direction: list, value=None, in_local_frame=False): + self.set_roll(0) + self.set_pitch(0) + if in_local_frame: + from metadrive.engine.engine_utils import get_engine + engine = get_engine() + direction = LVector3(*direction, 0.) + direction[1] *= -1 + ret = engine.worldNP.getRelativeVector(self.origin, direction) + direction = ret + speed = (norm(direction[0], direction[1]) + 1e-6) + if value is not None: + norm_ratio = value / speed + else: + norm_ratio = 1 + + self._body.setLinearVelocity( + LVector3(direction[0] * norm_ratio, direction[1] * norm_ratio, + self._body.getLinearVelocity()[-1]) + ) + self.standup() + + @property + def HEIGHT(self): + return self.config["height"] + + @property + def LENGTH(self): + return self.config["length"] + + @property + def WIDTH(self): + return self.config["width"] diff --git a/metadrive/component/vehicle/base_vehicle.py b/metadrive/component/vehicle/base_vehicle.py index 071a69d18..607289d69 100644 --- a/metadrive/component/vehicle/base_vehicle.py +++ b/metadrive/component/vehicle/base_vehicle.py @@ -145,6 +145,15 @@ def __init__( self.add_body(vehicle_chassis.getChassis()) self.system = vehicle_chassis self.chassis = self.origin + + if self.config["scale"] is not None: + w, l, h = self.config["scale"] + self.FRONT_WHEELBASE *= l + self.REAR_WHEELBASE *= l + self.LATERAL_TIRE_TO_CENTER *= w + self.TIRE_RADIUS *= h + self.CHASSIS_TO_WHEEL_AXIS *= h + self.wheels = self._create_wheel() # light experimental! @@ -626,18 +635,34 @@ def _create_vehicle_chassis(self): def _add_visualization(self): if self.render: - [path, scale, offset, HPR] = self.path - if path not in BaseVehicle.model_collection: + path, scale, offset, HPR = self.path + should_update = (path not in BaseVehicle.model_collection) or (self.config["scale"] is not None) + + if should_update: car_model = self.loader.loadModel(AssetLoader.file_path("models", path)) car_model.setTwoSided(False) - BaseVehicle.model_collection[path] = car_model + extra_offset_z = -self.TIRE_RADIUS - self.CHASSIS_TO_WHEEL_AXIS + if self.config['scale'] is not None: + offset = ( + offset[0] * self.config['scale'][0], offset[1] * self.config['scale'][1], + offset[2] * self.config['scale'][2] + ) + scale = (self.config['scale'][0], self.config['scale'][1], self.config['scale'][2]) + + # A quick workaround here. + # The model position is set to height/2 in ScenarioMapManager. + # Now we set this offset to -height/2, so that the model will be placed on the ground. + extra_offset_z = -self.config["height"] / 2 + car_model.setScale(scale) # model default, face to y car_model.setHpr(*HPR) - car_model.setPos(offset[0], offset[1], offset[-1]) - car_model.setZ(-self.TIRE_RADIUS - self.CHASSIS_TO_WHEEL_AXIS + offset[-1]) + car_model.setPos(offset[0], offset[1], offset[2] + extra_offset_z) + BaseVehicle.model_collection[path] = car_model + else: car_model = BaseVehicle.model_collection[path] + car_model.instanceTo(self.origin) if self.config["random_color"]: material = Material() @@ -679,7 +704,23 @@ def _add_wheel(self, pos: Vec3, radius: float, front: bool, left): wheel_model = self.loader.loadModel(model_path) wheel_model.setTwoSided(self.TIRE_TWO_SIDED) wheel_model.reparentTo(wheel_np) - wheel_model.set_scale(1 * self.TIRE_MODEL_CORRECT if left else -1 * self.TIRE_MODEL_CORRECT) + tire_scale = 1 * self.TIRE_MODEL_CORRECT if left else -1 * self.TIRE_MODEL_CORRECT + + if self.config['scale'] is not None: + tire_scale = ( + self.config['scale'][0] * tire_scale, self.config['scale'][1] * tire_scale, + self.config['scale'][2] * tire_scale + ) + + # A quick workaround here. + # The model position is set to height/2 in ScenarioMapManager. + # Now we set this offset to -height/2, so that the model will be placed on the ground. + # For the wheel, the bottom of it is not z=0, so we add two more terms to correct it. + extra_offset = -self.config["height"] / 2 + self.TIRE_RADIUS / self.config['scale'][ + 2] + self.CHASSIS_TO_WHEEL_AXIS / self.config['scale'][2] + wheel_model.setPos(0, 0, extra_offset) + + wheel_model.set_scale(tire_scale) wheel = self.system.createWheel() wheel.setNode(wheel_np.node()) wheel.setChassisConnectionPointCs(pos) diff --git a/metadrive/component/vehicle/vehicle_type.py b/metadrive/component/vehicle/vehicle_type.py index 98eab0a08..787a7485a 100644 --- a/metadrive/component/vehicle/vehicle_type.py +++ b/metadrive/component/vehicle/vehicle_type.py @@ -1,8 +1,12 @@ import platform -from metadrive.component.pg_space import ParameterSpace, VehicleParameterSpace +from panda3d.core import LineSegs, NodePath +from panda3d.core import Material, Vec3, LVecBase4 + +from metadrive.component.pg_space import VehicleParameterSpace, ParameterSpace from metadrive.component.vehicle.base_vehicle import BaseVehicle from metadrive.constants import Semantics +from metadrive.engine.asset_loader import AssetLoader class DefaultVehicle(BaseVehicle): @@ -16,19 +20,23 @@ class DefaultVehicle(BaseVehicle): LATERAL_TIRE_TO_CENTER = 0.815 FRONT_WHEELBASE = 1.05234 REAR_WHEELBASE = 1.4166 - path = ['ferra/vehicle.gltf', (1, 1, 1), (0, 0.075, 0.), (0, 0, 0)] + path = ('ferra/vehicle.gltf', (1, 1, 1), (0, 0.075, 0.), (0, 0, 0)) # asset path, scale, offset, HPR + + DEFAULT_LENGTH = 4.515 # meters + DEFAULT_HEIGHT = 1.19 # meters + DEFAULT_WIDTH = 1.852 # meters @property def LENGTH(self): - return 4.515 # meters + return self.DEFAULT_LENGTH @property def HEIGHT(self): - return 1.19 # meters + return self.DEFAULT_HEIGHT @property def WIDTH(self): - return 1.852 # meters + return self.DEFAULT_WIDTH # When using DefaultVehicle as traffic, please use this class. @@ -57,19 +65,23 @@ class XLVehicle(BaseVehicle): MASS = 1600 LIGHT_POSITION = (-0.75, 2.7, 0.2) SEMANTIC_LABEL = Semantics.TRUCK.label - path = ['truck/vehicle.gltf', (1, 1, 1), (0, 0.25, 0.04), (0, 0, 0)] + path = ('truck/vehicle.gltf', (1, 1, 1), (0, 0.25, 0.04), (0, 0, 0)) + + DEFAULT_LENGTH = 5.74 # meters + DEFAULT_HEIGHT = 2.8 # meters + DEFAULT_WIDTH = 2.3 # meters @property def LENGTH(self): - return 5.74 # meters + return self.DEFAULT_LENGTH @property def HEIGHT(self): - return 2.8 # meters + return self.DEFAULT_HEIGHT @property def WIDTH(self): - return 2.3 # meters + return self.DEFAULT_WIDTH class LVehicle(BaseVehicle): @@ -84,20 +96,23 @@ class LVehicle(BaseVehicle): TIRE_WIDTH = 0.35 MASS = 1300 LIGHT_POSITION = (-0.65, 2.13, 0.3) + DEFAULT_LENGTH = 4.87 # meters + DEFAULT_HEIGHT = 1.85 # meters + DEFAULT_WIDTH = 2.046 # meters path = ['lada/vehicle.gltf', (1.1, 1.1, 1.1), (0, -0.27, 0.07), (0, 0, 0)] @property def LENGTH(self): - return 4.87 # meters + return self.DEFAULT_LENGTH @property def HEIGHT(self): - return 1.85 # meters + return self.DEFAULT_HEIGHT @property def WIDTH(self): - return 2.046 # meters + return self.DEFAULT_WIDTH class MVehicle(BaseVehicle): @@ -112,20 +127,22 @@ class MVehicle(BaseVehicle): TIRE_WIDTH = 0.3 MASS = 1200 LIGHT_POSITION = (-0.67, 1.86, 0.22) - + DEFAULT_LENGTH = 4.6 # meters + DEFAULT_HEIGHT = 1.37 # meters + DEFAULT_WIDTH = 1.85 # meters path = ['130/vehicle.gltf', (1, 1, 1), (0, -0.05, 0.1), (0, 0, 0)] @property def LENGTH(self): - return 4.6 # meters + return self.DEFAULT_LENGTH @property def HEIGHT(self): - return 1.37 # meters + return self.DEFAULT_HEIGHT @property def WIDTH(self): - return 1.85 # meters + return self.DEFAULT_WIDTH class SVehicle(BaseVehicle): @@ -141,6 +158,9 @@ class SVehicle(BaseVehicle): TIRE_WIDTH = 0.25 MASS = 800 LIGHT_POSITION = (-0.57, 1.86, 0.23) + DEFAULT_LENGTH = 4.3 # meters + DEFAULT_HEIGHT = 1.7 # meters + DEFAULT_WIDTH = 1.7 # meters @property def path(self): @@ -154,15 +174,15 @@ def path(self): @property def LENGTH(self): - return 4.3 # meters + return self.DEFAULT_LENGTH @property def HEIGHT(self): - return 1.70 # meters + return self.DEFAULT_HEIGHT @property def WIDTH(self): - return 1.70 # meters + return self.DEFAULT_WIDTH class VaryingDynamicsVehicle(DefaultVehicle): @@ -208,25 +228,25 @@ def reset( if vehicle_config["length"] is not None and vehicle_config["length"] != self.LENGTH: should_force_reset = True if "max_engine_force" in vehicle_config and \ - vehicle_config["max_engine_force"] is not None and \ - vehicle_config["max_engine_force"] != self.config["max_engine_force"]: + vehicle_config["max_engine_force"] is not None and \ + vehicle_config["max_engine_force"] != self.config["max_engine_force"]: should_force_reset = True if "max_brake_force" in vehicle_config and \ - vehicle_config["max_brake_force"] is not None and \ - vehicle_config["max_brake_force"] != self.config["max_brake_force"]: + vehicle_config["max_brake_force"] is not None and \ + vehicle_config["max_brake_force"] != self.config["max_brake_force"]: should_force_reset = True if "wheel_friction" in vehicle_config and \ - vehicle_config["wheel_friction"] is not None and \ - vehicle_config["wheel_friction"] != self.config["wheel_friction"]: + vehicle_config["wheel_friction"] is not None and \ + vehicle_config["wheel_friction"] != self.config["wheel_friction"]: should_force_reset = True if "max_steering" in vehicle_config and \ - vehicle_config["max_steering"] is not None and \ - vehicle_config["max_steering"] != self.config["max_steering"]: + vehicle_config["max_steering"] is not None and \ + vehicle_config["max_steering"] != self.config["max_steering"]: self.max_steering = vehicle_config["max_steering"] should_force_reset = True if "mass" in vehicle_config and \ - vehicle_config["mass"] is not None and \ - vehicle_config["mass"] != self.config["mass"]: + vehicle_config["mass"] is not None and \ + vehicle_config["mass"] != self.config["mass"]: should_force_reset = True # def process_memory(): @@ -266,6 +286,138 @@ def reset( return ret +class VaryingDynamicsBoundingBoxVehicle(VaryingDynamicsVehicle): + def __init__( + self, vehicle_config: dict = None, name: str = None, random_seed=None, position=None, heading=None, **kwargs + ): + + # TODO(pzh): The above code is removed for now. How we get BUS label? + # vehicle_config has 'width' 'length' and 'height' + # if vehicle_config["width"] < 0.0: + # self.SEMANTIC_LABEL = Semantics.CAR.label + # else: + # self.SEMANTIC_LABEL = Semantics.BUS.label + + super(VaryingDynamicsBoundingBoxVehicle, self).__init__( + vehicle_config=vehicle_config, + name=name, + random_seed=random_seed, + position=position, + heading=heading, + **kwargs + ) + + def _add_visualization(self): + if self.render: + path, scale, offset, HPR = self.path + + # PZH: Note that we do not use model_collection as a buffer here. + # if path not in BaseVehicle.model_collection: + + # PZH: Load a box model and resize it to the vehicle size + car_model = AssetLoader.loader.loadModel(AssetLoader.file_path("models", "box.bam")) + + car_model.setTwoSided(False) + BaseVehicle.model_collection[path] = car_model + car_model.setScale((self.WIDTH, self.LENGTH, self.HEIGHT)) + # car_model.setZ(-self.TIRE_RADIUS - self.CHASSIS_TO_WHEEL_AXIS + self.HEIGHT / 2) + car_model.setZ(0) + # model default, face to y + car_model.setHpr(*HPR) + car_model.instanceTo(self.origin) + + show_contour = self.config["show_contour"] if "show_contour" in self.config else False + if show_contour: + # ========== Draw the contour of the bounding box ========== + # Draw the bottom of the car first + line_seg = LineSegs("bounding_box_contour1") + zoffset = car_model.getZ() + line_seg.setThickness(2) + line_color = [1.0, 0.0, 0.0] + out_offset = 0.02 + w = self.WIDTH / 2 + out_offset + l = self.LENGTH / 2 + out_offset + h = self.HEIGHT / 2 + out_offset + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(-w, l, h + zoffset) + line_seg.drawTo(-w, l, -h + zoffset) + line_seg.drawTo(w, l, -h + zoffset) + line_seg.drawTo(w, l, h + zoffset) + line_seg.drawTo(-w, l, -h + zoffset) + line_seg.moveTo(-w, l, h + zoffset) + line_seg.drawTo(w, l, -h + zoffset) + + line_seg.moveTo(w, -l, h + zoffset) + line_seg.drawTo(-w, -l, h + zoffset) + line_seg.drawTo(-w, -l, -h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_seg.drawTo(w, -l, h + zoffset) + line_seg.moveTo(-w, -l, 0 + zoffset) + line_seg.drawTo(w, -l, 0 + zoffset) + line_seg.moveTo(0, -l, h + zoffset) + line_seg.drawTo(0, -l, -h + zoffset) + + line_seg.moveTo(w, l, h + zoffset) + line_seg.drawTo(w, -l, h + zoffset) + line_seg.moveTo(-w, l, h + zoffset) + line_seg.drawTo(-w, -l, h + zoffset) + line_seg.moveTo(-w, l, -h + zoffset) + line_seg.drawTo(-w, -l, -h + zoffset) + line_seg.moveTo(w, l, -h + zoffset) + line_seg.drawTo(w, -l, -h + zoffset) + line_np = NodePath(line_seg.create(True)) + line_material = Material() + line_material.setBaseColor(LVecBase4(*line_color[:3], 1)) + line_np.setMaterial(line_material, True) + line_np.reparentTo(self.origin) + + if self.config["random_color"]: + material = Material() + material.setBaseColor( + ( + self.panda_color[0] * self.MATERIAL_COLOR_COEFF, + self.panda_color[1] * self.MATERIAL_COLOR_COEFF, + self.panda_color[2] * self.MATERIAL_COLOR_COEFF, 0. + ) + ) + material.setMetallic(self.MATERIAL_METAL_COEFF) + material.setSpecular(self.MATERIAL_SPECULAR_COLOR) + material.setRefractiveIndex(1.5) + material.setRoughness(self.MATERIAL_ROUGHNESS) + material.setShininess(self.MATERIAL_SHININESS) + material.setTwoside(False) + self.origin.setMaterial(material, True) + + def _add_wheel(self, pos: Vec3, radius: float, front: bool, left): + wheel_np = self.origin.attachNewNode("wheel") + self._node_path_list.append(wheel_np) + + # PZH: Skip the wheel model + # if self.render: + # model = 'right_tire_front.gltf' if front else 'right_tire_back.gltf' + # model_path = AssetLoader.file_path("models", os.path.dirname(self.path[0]), model) + # wheel_model = self.loader.loadModel(model_path) + # wheel_model.setTwoSided(self.TIRE_TWO_SIDED) + # wheel_model.reparentTo(wheel_np) + # wheel_model.set_scale(1 * self.TIRE_MODEL_CORRECT if left else -1 * self.TIRE_MODEL_CORRECT) + wheel = self.system.createWheel() + wheel.setNode(wheel_np.node()) + wheel.setChassisConnectionPointCs(pos) + wheel.setFrontWheel(front) + wheel.setWheelDirectionCs(Vec3(0, 0, -1)) + wheel.setWheelAxleCs(Vec3(1, 0, 0)) + + wheel.setWheelRadius(radius) + wheel.setMaxSuspensionTravelCm(self.SUSPENSION_LENGTH) + wheel.setSuspensionStiffness(self.SUSPENSION_STIFFNESS) + wheel.setWheelsDampingRelaxation(4.8) + wheel.setWheelsDampingCompression(1.2) + wheel_friction = self.config["wheel_friction"] if not self.config["no_wheel_friction"] else 0 + wheel.setFrictionSlip(wheel_friction) + wheel.setRollInfluence(0.5) + return wheel + + def random_vehicle_type(np_random, p=None): v_type = { "s": SVehicle, @@ -289,6 +441,7 @@ def random_vehicle_type(np_random, p=None): "default": DefaultVehicle, "static_default": StaticDefaultVehicle, "varying_dynamics": VaryingDynamicsVehicle, + "varying_dynamics_bounding_box": VaryingDynamicsBoundingBoxVehicle, "traffic_default": TrafficDefaultVehicle } diff --git a/metadrive/constants.py b/metadrive/constants.py index 00f5cc5bd..8c7602177 100644 --- a/metadrive/constants.py +++ b/metadrive/constants.py @@ -232,11 +232,13 @@ def collision_rules(cls): ] @classmethod - def set_collision_rule(cls, world: BulletWorld): + def set_collision_rule(cls, world: BulletWorld, disable_collision: bool = False): for rule in cls.collision_rules(): group_1 = int(math.log(rule[0].getWord(), 2)) group_2 = int(math.log(rule[1].getWord(), 2)) relation = rule[-1] + if disable_collision: + relation = False world.setGroupCollisionFlag(group_1, group_2, relation) @classmethod @@ -391,6 +393,12 @@ class Semantics: LANE_LINE = label_color("LANE_LINE", (255, 255, 255)) CROSSWALK = label_color("CROSSWALK", (55, 176, 189)) + # These color might be prettier? + # LANE_LINE = label_color("LANE_LINE", (128, 64, 128)) + # CROSSWALK = label_color("CROSSWALK", (128, 64, 128)) + + BUS = label_color("BUS", (0, 60, 100)) # PZH: I just randomly choose a color. + class MapTerrainSemanticColor: """ diff --git a/metadrive/engine/base_engine.py b/metadrive/engine/base_engine.py index 24b2b9f63..f5497fc2f 100644 --- a/metadrive/engine/base_engine.py +++ b/metadrive/engine/base_engine.py @@ -21,15 +21,18 @@ def generate_distinct_rgb_values(): - distinct_rgb_values = [] - step = 256 // 32 # 8 intervals for each RGB component (0-31, 32-63, ..., 224-255) + # Try to avoid (0,0,0) and (255,255,255) to avoid confusion with the background and other objects. + r = np.linspace(16, 256 - 16, 16).astype(int) + g = np.linspace(16, 256 - 16, 16).astype(int) + b = np.linspace(16, 256 - 16, 16).astype(int) - for r in range(step, 256, step): - for g in range(0, 256, step): - for b in range(0, 256, step): - distinct_rgb_values.append((round(r / 255, 5), round(g / 255, 5), round(b / 255, 5))) + # Create a meshgrid and reshape to get all combinations of r, g, b + rgbs = np.array(np.meshgrid(r, g, b)).T.reshape(-1, 3) - return distinct_rgb_values[:4096] # Return the first 4096 values + # Normalize the values to be between 0 and 1 + rgbs = rgbs / 255.0 + + return tuple(tuple(round(vv, 5) for vv in v) for v in rgbs) COLOR_SPACE = generate_distinct_rgb_values() @@ -151,8 +154,11 @@ def spawn_object(self, object_class, force_spawn=False, auto_fill_random_seed=Tr self._spawned_objects[obj.id] = obj color = self._pick_color(obj.id) if color == (-1, -1, -1): - print("FK!~") - exit() + raise ValueError( + "No color available for object: {} instance segment mask. We already used all {} colors...".format( + obj.id, BaseEngine.MAX_COLOR + ) + ) obj.attach_to_world(self.worldNP, self.physics_world) return obj @@ -186,11 +192,14 @@ def _clean_color(self, id): """ if id in self.id_c.keys(): - my_color = self.id_c[id] - BaseEngine.COLORS_OCCUPIED.remove(my_color) + my_color = self.id_c.pop(id) + if my_color in BaseEngine.COLORS_OCCUPIED: + BaseEngine.COLORS_OCCUPIED.remove(my_color) BaseEngine.COLORS_FREE.add(my_color) # print("After cleaning:,", len(BaseEngine.COLORS_OCCUPIED), len(BaseEngine.COLORS_FREE)) - self.id_c.pop(id) + # if id in self.id_c.keys(): + # self.id_c.pop(id) + assert my_color in self.c_id.keys() self.c_id.pop(my_color) def id_to_color(self, id): @@ -261,6 +270,7 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F else: raise ValueError("filter should be a list or a function") for id, obj in exclude_objects.items(): + self._clean_color(id) self._spawned_objects.pop(id) if id in self._object_tasks: self._object_tasks.pop(id) @@ -268,7 +278,7 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F policy = self._object_policies.pop(id) policy.destroy() if force_destroy_this_obj: - self._clean_color(obj.id) + #self._clean_color(obj.id) obj.destroy() else: obj.detach_from_world(self.physics_world) @@ -283,7 +293,7 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F if len(self._dying_objects[obj.class_name]) < self.global_config["num_buffering_objects"]: self._dying_objects[obj.class_name].append(obj) else: - self._clean_color(obj.id) + #self._clean_color(obj.id) obj.destroy() if self.global_config["record_episode"] and not self.replay_episode and record: self.record_manager.add_clear_info(obj) @@ -542,8 +552,8 @@ def register_manager(self, manager_name: str, manager): :param manager_name: name shouldn't exist in self._managers and not be same as any class attribute :param manager: subclass of BaseManager """ - assert manager_name not in self._managers, "Manager already exists in BaseEngine, Use update_manager() to " \ - "overwrite" + assert manager_name not in self._managers, "Manager {} already exists in BaseEngine, Use update_manager() to " \ + "overwrite".format(manager_name) assert not hasattr(self, manager_name), "Manager name can not be same as the attribute in BaseEngine" self._managers[manager_name] = manager setattr(self, manager_name, manager) diff --git a/metadrive/engine/core/engine_core.py b/metadrive/engine/core/engine_core.py index 1dc04d785..b6bffbd90 100644 --- a/metadrive/engine/core/engine_core.py +++ b/metadrive/engine/core/engine_core.py @@ -273,7 +273,9 @@ def __init__(self, global_config): self.common_filter = None # physics world - self.physics_world = PhysicsWorld(self.global_config["debug_static_world"]) + self.physics_world = PhysicsWorld( + self.global_config["debug_static_world"], disable_collision=self.global_config["disable_collision"] + ) # collision callback self.physics_world.dynamic_world.setContactAddedCallback(PythonCallbackObject(collision_callback)) diff --git a/metadrive/engine/core/main_camera.py b/metadrive/engine/core/main_camera.py index 46fc453bb..b9de65ede 100644 --- a/metadrive/engine/core/main_camera.py +++ b/metadrive/engine/core/main_camera.py @@ -623,13 +623,15 @@ def unmap(self, stream=None): self._cuda_buffer = check_cudart_err(cudart.cudaGraphicsUnmapResources(1, self.cuda_graphics_resource, stream)) return self - def get_image(self): - # The Tracked obj arg is only for compatibility - img = PNMImage() - self.engine.win.getScreenshot(img) - return img + # def get_image(self): + # # The Tracked obj arg is only for compatibility + # img = PNMImage() + # self.engine.win.getScreenshot(img) + # return img def save_image(self, tracked_obj, file_name="main_camera.png", **kwargs): # The Tracked obj arg is only for compatibility - img = self.get_image() + # img = self.get_image() + img = PNMImage() + self.engine.win.getScreenshot(img) img.write(file_name) diff --git a/metadrive/engine/core/physics_world.py b/metadrive/engine/core/physics_world.py index 933075c25..bfb8d529a 100644 --- a/metadrive/engine/core/physics_world.py +++ b/metadrive/engine/core/physics_world.py @@ -7,14 +7,14 @@ class PhysicsWorld: - def __init__(self, debug=False): + def __init__(self, debug=False, disable_collision=False): # a dynamic world, moving objects or objects which react to other objects should be placed here self.dynamic_world = BulletWorld() - CollisionGroup.set_collision_rule(self.dynamic_world) + CollisionGroup.set_collision_rule(self.dynamic_world, disable_collision=disable_collision) self.dynamic_world.setGravity(Vec3(0, 0, -9.81)) # set gravity # a static world which used to query position/overlap .etc. Don't implement doPhysics() in this world self.static_world = BulletWorld() if not debug else self.dynamic_world - CollisionGroup.set_collision_rule(self.static_world) + CollisionGroup.set_collision_rule(self.static_world, disable_collision=disable_collision) def report_bodies(self): dynamic_bodies = \ diff --git a/metadrive/envs/base_env.py b/metadrive/envs/base_env.py index 4102fe914..05174ec17 100644 --- a/metadrive/envs/base_env.py +++ b/metadrive/envs/base_env.py @@ -158,6 +158,7 @@ length=None, height=None, mass=None, + scale=None, # triplet (x, y, z) # Set the vehicle size only for pygame top-down renderer. It doesn't affect the physical size! top_down_width=None, @@ -211,6 +212,8 @@ preload_models=True, # model compression increasing the launch time disable_model_compression=True, + # Whether to disable the collision detection (useful for debugging / replay logged scenarios) + disable_collision=False, # ===== Terrain ===== # The size of the square map region, which is centered at [0, 0]. The map objects outside it are culled. diff --git a/metadrive/envs/scenario_env.py b/metadrive/envs/scenario_env.py index 7831742be..1ba065862 100644 --- a/metadrive/envs/scenario_env.py +++ b/metadrive/envs/scenario_env.py @@ -3,12 +3,12 @@ """ import numpy as np -from metadrive.component.navigation_module.edge_network_navigation import EdgeNetworkNavigation + from metadrive.component.navigation_module.trajectory_navigation import TrajectoryNavigation -from metadrive.constants import DEFAULT_AGENT from metadrive.constants import TerminationState from metadrive.engine.asset_loader import AssetLoader from metadrive.envs.base_env import BaseEnv +from metadrive.manager.scenario_agent_manager import ScenarioAgentManager from metadrive.manager.scenario_curriculum_manager import ScenarioCurriculumManager from metadrive.manager.scenario_data_manager import ScenarioDataManager from metadrive.manager.scenario_light_manager import ScenarioLightManager @@ -46,11 +46,11 @@ no_light=False, # no traffic light reactive_traffic=False, # turn on to enable idm traffic filter_overlapping_car=True, # If in one frame a traffic vehicle collides with ego car, it won't be created. - even_sample_vehicle_class=True, # to make the scene more diverse default_vehicle_in_traffic=False, skip_missing_light=True, static_traffic_object=True, show_sidewalk=False, + even_sample_vehicle_class=None, # Deprecated. # ===== Agent config ===== vehicle_config=dict( @@ -90,7 +90,8 @@ # ===== others ===== allowed_more_steps=None, # horizon, None=infinite - top_down_show_real_size=False + top_down_show_real_size=False, + use_bounding_box=False, # Set True to use a cube in visualization to represent every dynamic objects. ) @@ -115,6 +116,18 @@ def __init__(self, config=None): self.start_index = self.config["start_scenario_index"] self.num_scenarios = self.config["num_scenarios"] + def _post_process_config(self, config): + config = super(ScenarioEnv, self)._post_process_config(config) + if config["use_bounding_box"]: + config["vehicle_config"]["random_color"] = True + config["vehicle_config"]["vehicle_model"] = "varying_dynamics_bounding_box" + config["agent_configs"]["default_agent"]["use_special_color"] = True + config["agent_configs"]["default_agent"]["vehicle_model"] = "varying_dynamics_bounding_box" + return config + + def _get_agent_manager(self): + return ScenarioAgentManager(init_observations=self._get_observations()) + def setup_engine(self): super(ScenarioEnv, self).setup_engine() self.engine.register_manager("data_manager", ScenarioDataManager()) @@ -179,7 +192,7 @@ def msg(reason): done = True self.logger.info(msg("max step"), extra={"log_once": True}) elif self.config["allowed_more_steps"] and self.episode_lengths[vehicle_id] >= \ - self.engine.data_manager.current_scenario_length + self.config["allowed_more_steps"]: + self.engine.data_manager.current_scenario_length + self.config["allowed_more_steps"]: if self.config["truncate_as_terminate"]: done = True done_info[TerminationState.MAX_STEP] = True diff --git a/metadrive/examples/drive_in_real_env_with_bounding_box.py b/metadrive/examples/drive_in_real_env_with_bounding_box.py new file mode 100755 index 000000000..b26142a0a --- /dev/null +++ b/metadrive/examples/drive_in_real_env_with_bounding_box.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python +""" +This script demonstrates how to use the environment where traffic and road map are loaded from Waymo dataset. +""" +import argparse +import random + +from metadrive.constants import HELP_MESSAGE +from metadrive.engine.asset_loader import AssetLoader +from metadrive.envs.scenario_env import ScenarioEnv + +RENDER_MESSAGE = { + "Quit": "ESC", + "Switch perspective": "Q or B", + "Reset Episode": "R", + "Keyboard Control": "W,A,S,D", +} + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--reactive_traffic", action="store_true") + parser.add_argument("--top_down", "--topdown", action="store_true") + parser.add_argument("--waymo", action="store_true") + args = parser.parse_args() + extra_args = dict(film_size=(2000, 2000)) if args.top_down else {} + asset_path = AssetLoader.asset_path + use_waymo = args.waymo + print(HELP_MESSAGE) + try: + env = ScenarioEnv( + { + "manual_control": True, + "sequential_seed": True, + "reactive_traffic": True if args.reactive_traffic else False, + "use_render": True if not args.top_down else False, + "data_directory": AssetLoader.file_path( + asset_path, "waymo" if use_waymo else "nuscenes", unix_style=False + ), + "num_scenarios": 3 if use_waymo else 10, + "debug": True, + "use_bounding_box": True, + "vehicle_config": { + "show_line_to_dest": True, + "show_line_to_navi_mark": True, + }, + "disable_collision": True, + } + ) + o, _ = env.reset() + + for i in range(1, 100000): + o, r, tm, tc, info = env.step([1.0, 0.]) + env.render( + mode="top_down" if args.top_down else None, + text=None if args.top_down else RENDER_MESSAGE, + **extra_args + ) + if tm or tc: + env.reset() + finally: + env.close() diff --git a/metadrive/manager/scenario_agent_manager.py b/metadrive/manager/scenario_agent_manager.py new file mode 100644 index 000000000..00e8cd9c3 --- /dev/null +++ b/metadrive/manager/scenario_agent_manager.py @@ -0,0 +1,8 @@ +from metadrive.manager.agent_manager import VehicleAgentManager + + +class ScenarioAgentManager(VehicleAgentManager): + """ + PZH: I should note that the ego car's information is filled by scenario_map_manager. + So we don't have to do anything here. + """ diff --git a/metadrive/manager/scenario_map_manager.py b/metadrive/manager/scenario_map_manager.py index 5a4fd51c0..b46dcd4c4 100644 --- a/metadrive/manager/scenario_map_manager.py +++ b/metadrive/manager/scenario_map_manager.py @@ -4,6 +4,9 @@ from metadrive.constants import DEFAULT_AGENT from metadrive.manager.base_manager import BaseManager from metadrive.scenario.parse_object_state import parse_full_trajectory, parse_object_state, get_idm_route +from metadrive.engine.logger import get_logger, set_log_level + +logger = get_logger() class ScenarioMapManager(BaseManager): @@ -53,9 +56,28 @@ def update_route(self): sdc_traj = parse_full_trajectory(sdc_track) - init_state = parse_object_state(sdc_track, 0, check_last_state=False) + init_state = parse_object_state(sdc_track, 0, check_last_state=False, include_z_position=True) + + # PZH: There is a wierd bug in the nuscene's source data, the width and length of the object is not consistent. + # Maybe this should be handle in ScenarioNet. But for now, we have to handle it here. + # As a workaround, we swap the width and length if the width is larger than length. + if data["version"].startswith("nuscenesv1.0") or data["metadata"]["dataset"] == "nuscenes": + if init_state["width"] > init_state["length"]: + init_state["width"], init_state["length"] = init_state["length"], init_state["width"] + + if max(init_state["width"], init_state["length"]) > 2 and (init_state["width"] > init_state["length"]): + logger.warning( + "The width of the object {} is larger than length {}. Are you sure?".format( + init_state["width"], init_state["length"] + ) + ) + last_state = parse_object_state(sdc_track, -1, check_last_state=True) init_position = init_state["position"] + + # Add a fake Z axis so that the object will not fall from the sky. + init_position[-1] = 0 + init_yaw = init_state["heading"] last_position = last_state["position"] last_yaw = last_state["heading"] @@ -69,7 +91,11 @@ def update_route(self): dict( agent_configs={ DEFAULT_AGENT: dict( - spawn_position_heading=(init_position, init_yaw), spawn_velocity=init_state["velocity"] + spawn_position_heading=(list(init_position), init_yaw), + spawn_velocity=init_state["velocity"], + width=init_state["width"], + length=init_state["length"], + height=init_state["height"], ) } ) diff --git a/metadrive/manager/scenario_traffic_manager.py b/metadrive/manager/scenario_traffic_manager.py index ac25d53ef..e43c5f831 100644 --- a/metadrive/manager/scenario_traffic_manager.py +++ b/metadrive/manager/scenario_traffic_manager.py @@ -1,15 +1,15 @@ import copy -from metadrive.engine.logger import get_logger -from metadrive.utils.math import norm + import numpy as np from metadrive.component.static_object.traffic_object import TrafficCone, TrafficBarrier -from metadrive.component.traffic_participants.cyclist import Cyclist -from metadrive.component.traffic_participants.pedestrian import Pedestrian +from metadrive.component.traffic_participants.cyclist import Cyclist, CyclistBoundingBox +from metadrive.component.traffic_participants.pedestrian import Pedestrian, PedestrianBoundingBox from metadrive.component.vehicle.base_vehicle import BaseVehicle -from metadrive.component.vehicle.vehicle_type import SVehicle, LVehicle, MVehicle, XLVehicle, \ - TrafficDefaultVehicle +from metadrive.component.vehicle.vehicle_type import LVehicle, MVehicle, XLVehicle, \ + VaryingDynamicsBoundingBoxVehicle, SVehicle, DefaultVehicle from metadrive.constants import DEFAULT_AGENT +from metadrive.engine.logger import get_logger from metadrive.manager.base_manager import BaseManager from metadrive.policy.idm_policy import TrajectoryIDMPolicy from metadrive.policy.replay_policy import ReplayEgoCarPolicy @@ -17,6 +17,7 @@ from metadrive.scenario.parse_object_state import parse_object_state, get_idm_route, get_max_valid_indicis from metadrive.scenario.scenario_description import ScenarioDescription as SD from metadrive.type import MetaDriveType +from metadrive.utils.math import norm from metadrive.utils.math import wrap_to_pi logger = get_logger() @@ -56,7 +57,10 @@ def __init__(self): self._obj_to_clean_this_frame = [] # some flags - self.even_sample_v = self.engine.global_config["even_sample_vehicle_class"] + self.even_sample_v = self.engine.global_config.get("even_sample_vehicle_class", None) + if self.even_sample_v is not None: + raise DeprecationWarning("even_sample_vehicle_class is deprecated!") + self.need_default_vehicle = self.engine.global_config["default_vehicle_in_traffic"] self.is_ego_vehicle_replay = self.engine.global_config["agent_policy"] == ReplayEgoCarPolicy self._filter_overlapping_car = self.engine.global_config["filter_overlapping_car"] @@ -78,7 +82,7 @@ def before_step(self, *args, **kwargs): def before_reset(self): super(ScenarioTrafficManager, self).before_reset() self._obj_to_clean_this_frame = [] - reset_vehicle_type_count(self.np_random) + # reset_vehicle_type_count(self.np_random) def after_reset(self): self._scenario_id_to_obj_id = {} @@ -169,7 +173,10 @@ def vehicles(self): return list(self.engine.get_objects(filter=lambda o: isinstance(o, BaseVehicle)).values()) def spawn_vehicle(self, v_id, track): - state = parse_object_state(track, self.episode_step) + state = parse_object_state(track, self.episode_step, include_z_position=False) + use_bounding_box = ( + self.engine.global_config["vehicle_config"]["vehicle_model"] == "varying_dynamics_bounding_box" + ) # for each vehicle, we would like to know if it is static if v_id not in self._static_car_id and v_id not in self._moving_car_id: @@ -184,21 +191,31 @@ def spawn_vehicle(self, v_id, track): # if collision don't generate, unless ego car is in replay mode ego_pos = self.ego_vehicle.position - heading_dist, side_dist = self.ego_vehicle.convert_to_local_coordinates(state["position"], ego_pos) + heading_dist, side_dist = self.ego_vehicle.convert_to_local_coordinates(state["position"][:2], ego_pos) if not self.is_ego_vehicle_replay and self._filter_overlapping_car and \ abs(heading_dist) < self.GENERATION_FORWARD_CONSTRAINT and \ abs(side_dist) < self.GENERATION_SIDE_CONSTRAINT: return # create vehicle - if state["vehicle_class"]: + if state["vehicle_class"] and not use_bounding_box: vehicle_class = state["vehicle_class"] else: vehicle_class = get_vehicle_type( - float(state["length"]), None if self.even_sample_v else self.np_random, self.need_default_vehicle + float(state["length"]), self.need_default_vehicle, use_bounding_box=use_bounding_box ) + # print("vehicle_class: ", vehicle_class) obj_name = v_id if self.engine.global_config["force_reuse_object_name"] else None v_cfg = copy.copy(self._traffic_v_config) + + v_cfg["width"] = state["width"] + v_cfg["length"] = state["length"] + v_cfg["height"] = state["height"] + v_cfg["scale"] = ( + v_cfg["width"] / vehicle_class.DEFAULT_WIDTH, v_cfg["length"] / vehicle_class.DEFAULT_LENGTH, + v_cfg["height"] / vehicle_class.DEFAULT_HEIGHT + ) + if self.engine.global_config["top_down_show_real_size"]: v_cfg["top_down_length"] = track["state"]["length"][self.episode_step] v_cfg["top_down_width"] = track["state"]["width"][self.episode_step] @@ -207,8 +224,16 @@ def spawn_vehicle(self, v_id, track): "Scenario ID: {}. The top_down size of vehicle {} is weird: " "{}".format(self.engine.current_seed, v_id, [v_cfg["length"], v_cfg["width"]]) ) + + position = list(state["position"]) + + # Add z to make it stick to the ground: + assert len(position) == 2 + if use_bounding_box: + position.append(state['height'] / 2) + v = self.spawn_object( - vehicle_class, position=state["position"], heading=state["heading"], vehicle_config=v_cfg, name=obj_name + vehicle_class, position=position, heading=state["heading"], vehicle_config=v_cfg, name=obj_name ) self._scenario_id_to_obj_id[v_id] = v.name self._obj_id_to_scenario_id[v.name] = v_id @@ -236,15 +261,27 @@ def spawn_vehicle(self, v_id, track): self.idm_policy_count += 1 def spawn_pedestrian(self, scenario_id, track): - state = parse_object_state(track, self.episode_step) + state = parse_object_state(track, self.episode_step, include_z_position=False) if not state["valid"]: return obj_name = scenario_id if self.engine.global_config["force_reuse_object_name"] else None + if self.global_config["use_bounding_box"]: + cls = PedestrianBoundingBox + force_spawn = True + else: + cls = Pedestrian + force_spawn = False + + position = list(state["position"]) obj = self.spawn_object( - Pedestrian, + cls, name=obj_name, - position=state["position"], + position=position, heading_theta=state["heading"], + width=state["width"], + length=state["length"], + height=state["height"], + force_spawn=force_spawn ) self._scenario_id_to_obj_id[scenario_id] = obj.name self._obj_id_to_scenario_id[obj.name] = scenario_id @@ -252,15 +289,27 @@ def spawn_pedestrian(self, scenario_id, track): policy.act() def spawn_cyclist(self, scenario_id, track): - state = parse_object_state(track, self.episode_step) + state = parse_object_state(track, self.episode_step, include_z_position=False) if not state["valid"]: return obj_name = scenario_id if self.engine.global_config["force_reuse_object_name"] else None + if self.global_config["use_bounding_box"]: + cls = CyclistBoundingBox + force_spawn = True + else: + cls = Cyclist + force_spawn = False + + position = list(state["position"]) obj = self.spawn_object( - Cyclist, + cls, name=obj_name, - position=state["position"], + position=position, heading_theta=state["heading"], + width=state["width"], + length=state["length"], + height=state["height"], + force_spawn=force_spawn ) self._scenario_id_to_obj_id[scenario_id] = obj.name self._obj_id_to_scenario_id[obj.name] = scenario_id @@ -303,7 +352,7 @@ def ego_vehicle(self): def is_static_object(self, obj_id): return isinstance(self.spawned_objects[obj_id], TrafficBarrier) \ - or isinstance(self.spawned_objects[obj_id], TrafficCone) + or isinstance(self.spawned_objects[obj_id], TrafficCone) @property def obj_id_to_scenario_id(self): @@ -333,38 +382,27 @@ def get_traffic_v_config(): return v_config -type_count = [0 for i in range(3)] +# type_count = [0 for i in range(3)] -def get_vehicle_type(length, np_random=None, need_default_vehicle=False): - if np_random is not None: - if length <= 4: - return SVehicle - elif length <= 5.5: - return [LVehicle, SVehicle, MVehicle][np_random.randint(3)] - else: - return [LVehicle, XLVehicle][np_random.randint(2)] +def get_vehicle_type(length, need_default_vehicle=False, use_bounding_box=False): + if use_bounding_box: + return VaryingDynamicsBoundingBoxVehicle + if need_default_vehicle: + return DefaultVehicle + if length <= 4: + return SVehicle + elif length <= 5.2: + return MVehicle + elif length <= 6.2: + return LVehicle else: - global type_count - # evenly sample - if length <= 4: - return SVehicle - elif length <= 5.5: - type_count[1] += 1 - vs = [LVehicle, MVehicle, SVehicle] - # vs = [SVehicle, LVehicle, MVehicle] - if need_default_vehicle: - vs.append(TrafficDefaultVehicle) - return vs[type_count[1] % len(vs)] - else: - type_count[2] += 1 - vs = [LVehicle, XLVehicle] - return vs[type_count[2] % len(vs)] + return XLVehicle -def reset_vehicle_type_count(np_random=None): - global type_count - if np_random is None: - type_count = [0 for i in range(3)] - else: - type_count = [np_random.randint(100) for i in range(3)] +# def reset_vehicle_type_count(np_random=None): +# global type_count +# if np_random is None: +# type_count = [0 for i in range(3)] +# else: +# type_count = [np_random.randint(100) for i in range(3)] diff --git a/metadrive/policy/replay_policy.py b/metadrive/policy/replay_policy.py index 80f54d8a7..28ccce036 100644 --- a/metadrive/policy/replay_policy.py +++ b/metadrive/policy/replay_policy.py @@ -63,6 +63,7 @@ def act(self, *args, **kwargs): self.control_object.set_velocity(info["velocity"], in_local_frame=self._velocity_local_frame) self.control_object.set_heading_theta(info["heading"]) self.control_object.set_angular_velocity(info["angular_velocity"]) + self.control_object.set_static(True) return None # Return None action so the base vehicle will not overwrite the steering & throttle diff --git a/metadrive/scenario/parse_object_state.py b/metadrive/scenario/parse_object_state.py index 6675ae03e..8078d2fb7 100644 --- a/metadrive/scenario/parse_object_state.py +++ b/metadrive/scenario/parse_object_state.py @@ -25,7 +25,7 @@ def get_idm_route(traj_points, width=2): return traj -def parse_object_state(object_dict, time_idx, check_last_state=False, sim_time_interval=0.1): +def parse_object_state(object_dict, time_idx, check_last_state=False, sim_time_interval=0.1, include_z_position=False): """ Parse object state of one time step """ @@ -47,7 +47,11 @@ def parse_object_state(object_dict, time_idx, check_last_state=False, sim_time_i ret = {k: v[time_idx] for k, v in states.items()} - ret["position"] = states["position"][time_idx, :2] + if include_z_position: + ret["position"] = states["position"][time_idx] + else: + ret["position"] = states["position"][time_idx, :2] + ret["velocity"] = states["velocity"][time_idx] ret["heading_theta"] = states["heading"][time_idx] diff --git a/metadrive/scenario/utils.py b/metadrive/scenario/utils.py index 346cbd297..c6b2242c9 100644 --- a/metadrive/scenario/utils.py +++ b/metadrive/scenario/utils.py @@ -166,8 +166,9 @@ def convert_recorded_scenario_exported(record_episode, scenario_log_interval=0.1 light_manager_name = find_light_manager_name(record_episode["manager_metadata"]) data_manager_name = find_data_manager_name(record_episode["manager_metadata"]) - tracks = { - k: dict( + tracks = {} + for k in list(all_objs): + tracks[k] = dict( type=MetaDriveType.UNSET, state=dict( position=np.zeros(shape=(episode_len, 3)), @@ -182,8 +183,6 @@ def convert_recorded_scenario_exported(record_episode, scenario_log_interval=0.1 ), metadata=dict(track_length=episode_len, type=MetaDriveType.UNSET, object_id=k, original_id=k) ) - for k in list(all_objs) - } all_lights = set() if light_manager_name is not None: @@ -245,6 +244,22 @@ def convert_recorded_scenario_exported(record_episode, scenario_log_interval=0.1 tracks[id]["state"]["position"][frame_idx] = state["position"] tracks[id]["state"]["heading"][frame_idx] = state["heading_theta"] tracks[id]["state"]["velocity"][frame_idx] = state["velocity"] + if "width" in state: + if "width" not in tracks[id]["state"]: + tracks[id]["state"]["width"] = np.zeros(shape=(episode_len, 1)) + tracks[id]["state"]["width"][frame_idx] = state["width"] + if "length" in state: + if "length" not in tracks[id]["state"]: + tracks[id]["state"]["length"] = np.zeros(shape=(episode_len, 1)) + tracks[id]["state"]["length"][frame_idx] = state["length"] + if "height" in state: + if "height" not in tracks[id]["state"]: + tracks[id]["state"]["height"] = np.zeros(shape=(episode_len, 1)) + tracks[id]["state"]["height"][frame_idx] = state["height"] + if "radius" in state: + if "radius" not in tracks[id]["state"]: + tracks[id]["state"]["radius"] = np.zeros(shape=(episode_len, 1)) + tracks[id]["state"]["radius"][frame_idx] = state["radius"] tracks[id]["state"]["valid"][frame_idx] = 1 if "throttle_brake" in state: diff --git a/metadrive/tests/test_component/test_lane_line_detector.py b/metadrive/tests/test_component/test_lane_line_detector.py index 3fa091eea..9563020dc 100644 --- a/metadrive/tests/test_component/test_lane_line_detector.py +++ b/metadrive/tests/test_component/test_lane_line_detector.py @@ -484,131 +484,131 @@ def test_pg_map(render=False): ] nuscenes_gt_2 = [ 1.0, - 0.7808130979537964, - 0.4014144241809845, - 0.497048556804657, - 0.21977420151233673, - 0.1833307445049286, - 0.2741309106349945, - 0.24397589266300201, - 0.22296150028705597, - 0.20816059410572052, - 0.19839195907115936, - 0.19227449595928192, - 0.18942095339298248, - 0.1900409460067749, - 0.19369441270828247, - 0.20063060522079468, - 0.1172584742307663, - 0.12590330839157104, - 0.13827820122241974, - 0.15608078241348267, - 0.18244224786758423, - 0.2237718254327774, - 0.2953447997570038, - 0.4444361925125122, - 0.9753275513648987, - 0.11675211787223816, - 0.04168549180030823, - 0.038301981985569, - 0.035907551646232605, - 0.034303370863199234, - 0.033339571207761765, - 0.03293364867568016, - 0.03305203467607498, - 0.035654980689287186, - 0.0451875701546669, - 0.06304079294204712, - 1.0, - 0.9279412627220154, - 0.47582244873046875, - 0.3268144130706787, - 0.25280630588531494, - 0.20891083776950836, - 0.18038643896579742, - 0.16094817221164703, - 0.1473899632692337, - 0.1379445493221283, - 0.1312660276889801, - 0.1269902139902115, - 0.12537911534309387, - 0.1267208307981491, - 0.3988616168498993, - 0.39293304085731506, + 0.779081404209137, + 0.4008123576641083, + 0.4965006411075592, + 0.2193922996520996, + 0.1830165535211563, + 0.27383944392204285, + 0.243718221783638, + 0.22272785007953644, + 0.20793406665325165, + 0.19818702340126038, + 0.19207695126533508, + 0.1892271488904953, + 0.18984796106815338, + 0.1934981793165207, + 0.20042887330055237, + 0.11704640090465546, + 0.1256762593984604, + 0.1380302459001541, + 0.15554210543632507, + 0.1821221262216568, + 0.22327348589897156, + 0.294590562582016, + 0.4437068700790405, + 0.9753552079200745, + 0.12470095604658127, + 0.04193984344601631, + 0.03848263621330261, + 0.03607578203082085, + 0.034464310854673386, + 0.03349597752094269, + 0.03308868035674095, + 0.033208273351192474, + 0.03570190817117691, + 0.04525022581219673, + 0.06312886625528336, + 1.0, + 0.9281278848648071, + 0.4757794439792633, + 0.326820969581604, + 0.2528233230113983, + 0.2089240849018097, + 0.1803993582725525, + 0.16096088290214539, + 0.1474027931690216, + 0.1379578709602356, + 0.13127265870571136, + 0.12699860334396362, + 0.12540613114833832, + 0.1267486810684204, + 0.39884063601493835, + 0.39179351925849915, 0.5, 0.5, 0.5, - 0.00011851061572087929, - 1.0, - 0.7808130979537964, - 0.4014144241809845, - 0.497048556804657, - 0.21977420151233673, - 0.1833307445049286, - 0.2741309106349945, - 0.24397589266300201, - 0.22296150028705597, - 0.20816059410572052, - 0.19839195907115936, - 0.19227449595928192, - 0.18942095339298248, - 0.1900409460067749, - 0.19369441270828247, - 0.20063060522079468, - 0.1172584742307663, - 0.12590330839157104, - 0.13827820122241974, - 0.15608078241348267, - 0.18244224786758423, - 0.2237718254327774, - 0.2953447997570038, - 0.4444361925125122, - 0.9753275513648987, - 0.11675211787223816, - 0.04168549180030823, - 0.038301981985569, - 0.035907551646232605, - 0.034303370863199234, - 0.033339571207761765, - 0.03293364867568016, - 0.03305203467607498, - 0.035654980689287186, - 0.0451875701546669, - 0.06304079294204712, - 1.0, - 0.9279412627220154, - 0.47582244873046875, - 0.3268144130706787, - 0.25280630588531494, - 0.20891083776950836, - 0.18038643896579742, - 0.16094817221164703, - 0.1473899632692337, - 0.1379445493221283, - 0.1312660276889801, - 0.1269902139902115, - 0.12537911534309387, - 0.1267208307981491, - 0.4396708011627197, - 0.4396708011627197, - 0.4618304371833801, - 0.4618304371833801, - 0.483804315328598, - 0.483804315328598, - 0.505452036857605, - 0.505452036857605, - 0.5266749858856201, - 0.5266749858856201, - 0.547338604927063, - 0.547338604927063, - 0.5673576593399048, - 0.5673576593399048, - 0.5866857767105103, - 0.5866857767105103, - 0.6053164601325989, - 0.6053164601325989, + 2.9294713385752402e-05, + 1.0, + 0.779081404209137, + 0.176986426115036, + 0.4965006411075592, + 0.2193922996520996, + 0.1830165535211563, + 0.27383944392204285, + 0.243718221783638, + 0.22272785007953644, + 0.20793406665325165, + 0.19818702340126038, + 0.19207695126533508, + 0.1892271488904953, + 0.18984796106815338, + 0.1934981793165207, + 0.20042887330055237, + 0.11704640090465546, + 0.1256762593984604, + 0.1380302459001541, + 0.15554210543632507, + 0.1821221262216568, + 0.22327348589897156, + 0.294590562582016, + 0.4437068700790405, + 0.9753552079200745, + 0.12470095604658127, + 0.04193984344601631, + 0.03848263621330261, + 0.03607578203082085, + 0.034464310854673386, + 0.03349597752094269, + 0.03308868035674095, + 0.033208273351192474, + 0.03570190817117691, + 0.04525022581219673, + 0.06312886625528336, + 1.0, + 0.9281278848648071, + 0.4757794439792633, + 0.326820969581604, + 0.2528233230113983, + 0.2089240849018097, + 0.1803993582725525, + 0.16096088290214539, + 0.1474027931690216, + 0.1379578709602356, + 0.13127265870571136, + 0.12699860334396362, + 0.12540613114833832, + 0.1267486810684204, + 0.4397680163383484, + 0.4397680163383484, + 0.46192362904548645, + 0.46192362904548645, + 0.48389288783073425, + 0.48389288783073425, + 0.5055355429649353, + 0.5055355429649353, + 0.5267529487609863, + 0.5267529487609863, + 0.5474106073379517, + 0.5474106073379517, + 0.5674235820770264, + 0.5674235820770264, + 0.5867453813552856, + 0.5867453813552856, + 0.6053695678710938, + 0.6053695678710938, 0.0, - 0.46758297085762024, + 0.46757614612579346, 0.0, 0.0, 1.0, @@ -679,5 +679,5 @@ def test_nuscenes(render=False): if __name__ == '__main__': - # test_nuscenes(True) - test_pg_map(False) + test_nuscenes(False) + # test_pg_map(False) diff --git a/metadrive/tests/test_export_record_scenario/test_export_scenario.py b/metadrive/tests/test_export_record_scenario/test_export_scenario.py index b8bfc531f..59077e72d 100644 --- a/metadrive/tests/test_export_record_scenario/test_export_scenario.py +++ b/metadrive/tests/test_export_record_scenario/test_export_scenario.py @@ -99,5 +99,5 @@ def test_export_waymo_scenario(num_scenarios=3, render_export_env=False, render_ if __name__ == "__main__": - test_export_metadrive_scenario(render_export_env=False, render_load_env=False) - # test_export_waymo_scenario(num_scenarios=3, render_export_env=False, render_load_env=False) + # test_export_metadrive_scenario(render_export_env=False, render_load_env=False) + test_export_waymo_scenario(num_scenarios=1, render_export_env=False, render_load_env=False) diff --git a/metadrive/tests/test_export_record_scenario/test_export_scenario_consistency_test.py b/metadrive/tests/test_export_record_scenario/test_export_scenario_consistency_test.py index 316c4d757..dac052a4b 100644 --- a/metadrive/tests/test_export_record_scenario/test_export_scenario_consistency_test.py +++ b/metadrive/tests/test_export_record_scenario/test_export_scenario_consistency_test.py @@ -426,9 +426,9 @@ def test_nuscenes_export_and_original_consistency(num_scenarios=7, render_export if __name__ == "__main__": # test_export_metadrive_scenario_reproduction(num_scenarios=10) - test_export_metadrive_scenario_easy(render_export_env=False, render_load_env=False) + # test_export_metadrive_scenario_easy(render_export_env=False, render_load_env=False) # test_export_metadrive_scenario_hard(num_scenarios=3, render_export_env=True, render_load_env=True) # test_export_waymo_scenario(num_scenarios=3, render_export_env=False, render_load_env=False) # test_waymo_export_and_original_consistency(num_scenarios=3, render_export_env=False) - # test_export_nuscenes_scenario(num_scenarios=2, render_export_env=False, render_load_env=False) + test_export_nuscenes_scenario(num_scenarios=1, render_export_env=False, render_load_env=False) # test_nuscenes_export_and_original_consistency() diff --git a/metadrive/tests/test_sensors/test_instance_cam.py b/metadrive/tests/test_sensors/test_instance_cam.py index 5ae806f6e..e5ac8b0e5 100644 --- a/metadrive/tests/test_sensors/test_instance_cam.py +++ b/metadrive/tests/test_sensors/test_instance_cam.py @@ -42,7 +42,7 @@ def test_instance_cam(config, render=False): env.reset() base_free = len(env.engine.COLORS_FREE) base_occupied = len(env.engine.COLORS_OCCUPIED) - assert base_free + base_occupied == 4096 + assert base_free + base_occupied == env.engine.MAX_COLOR import cv2 import time start = time.time() @@ -63,14 +63,20 @@ def test_instance_cam(config, render=False): color = unique_color.tolist() color = ( round(color[2], 5), round(color[1], 5), round(color[0], 5) - ) #In engine, we use 5-diigt float for keys + ) #In engine, we use 5-digit float for keys + if color not in env.engine.COLORS_OCCUPIED: + import matplotlib.pyplot as plt + plt.imshow(o["image"][..., -1]) + plt.show() + print("Unique colors:", unique_colors) + print("Occupied colors:", env.engine.COLORS_OCCUPIED) assert color in env.engine.COLORS_OCCUPIED assert color not in env.engine.COLORS_FREE assert color in env.engine.c_id.keys() assert env.engine.id_c[env.engine.c_id[color]] == color #Making sure the color-id is a bijection assert len(env.engine.c_id.keys()) == len(env.engine.COLORS_OCCUPIED) assert len(env.engine.id_c.keys()) == len(env.engine.COLORS_OCCUPIED) - assert len(env.engine.COLORS_FREE) + len(env.engine.COLORS_OCCUPIED) == 4096 + assert len(env.engine.COLORS_FREE) + len(env.engine.COLORS_OCCUPIED) == env.engine.MAX_COLOR #Making sure every object in the engine(not necessarily observable) have corresponding color for id, object in env.engine.get_objects().items(): assert id in env.engine.id_c.keys() @@ -84,4 +90,4 @@ def test_instance_cam(config, render=False): if __name__ == '__main__': test_instance_cam(config=blackbox_test_configs["small"], render=True) - my_dict = {(0, 0, 0): "Hello, World!"} + # my_dict = {(0, 0, 0): "Hello, World!"}