duburcqa / jiminy

Jiminy: a fast and portable Python/C++ simulator of poly-articulated robots with OpenAI Gym interface for reinforcement learning
https://duburcqa.github.io/jiminy
MIT License
219 stars 26 forks source link

[viewer] Add Open3d backend #602

Closed duburcqa closed 1 month ago

duburcqa commented 1 year ago

Open3d has the potential to replace both meshcat and panda3d. It seems to surpass them on many aspects. First, it has native web and notebook (the official way, without any hack) in addition to classical standalone window. Next, it seems to better support non-nvidia GPU machines. Besides, it is well maintained and not bloated by physics engine aspect such as Panda3d. It would be interesting to compare the performance wrt panda3d, especially during offscreen rendering. Here is a snippet for loading meshes (stl or dae) in Open3d 0.17:

python3 -m pip install open3d-cpu gym_jiminy_zoo pycollada
import os
import warnings
from pkg_resources import resource_filename
from typing import Optional

import numpy as np

import trimesh
import trimesh.viewer
import open3d as o3d
import open3d.visualization as vis

import gym
import jiminy_py
import hppfcl
import pinocchio as pin

# vis.webrtc_server.enable_webrtc()

def make_capsule(radius: float,
                 length: float,
                 num_segments: int = 16,
                 num_rings: int = 16
                 ) -> o3d.t.geometry.TriangleMesh:
    vertices, normals = [], []
    for u in np.linspace(0, np.pi, num_rings):
        for v in np.linspace(0, 2 * np.pi, num_segments):
            x, y, z = np.cos(v) * np.sin(u), np.sin(v) * np.sin(u), np.cos(u)
            offset = np.sign(z) * 0.5 * length
            vertices.append((x * radius, y * radius, z * radius + offset))
            normals.append((x, y, z))

    faces = []
    for i in range(num_rings - 1):
        for j in range(num_segments - 1):
            r0 = i * num_segments + j
            r1 = r0 + num_segments
            if i < num_rings - 2:
                faces.append((r0, r1, r1 + 1))
            if i > 0:
                faces.append((r0, r1 + 1, r0 + 1))

    mesh = o3d.t.geometry.TriangleMesh()
    mesh.vertex.positions = o3d.core.Tensor(vertices, o3d.core.float32)
    mesh.vertex.normals = o3d.core.Tensor(normals, o3d.core.float32)
    mesh.triangle.indices = o3d.core.Tensor(faces, o3d.core.int32)

    return mesh

def load_mesh_file(mesh_path: str,
                   scale: Optional[np.ndarray] = None
                   ) -> o3d.t.geometry.TriangleMesh:
    mesh = trimesh.load(mesh_path, force='mesh')

    mesh_o3d = o3d.t.geometry.TriangleMesh()
    mesh_o3d.vertex.positions = o3d.core.Tensor(mesh.vertices, o3d.core.float32)
    mesh_o3d.triangle.indices = o3d.core.Tensor(mesh.faces, o3d.core.int32)
    mesh_o3d.compute_vertex_normals()

    mesh_o3d.material.set_default_properties()
    mesh_o3d.material.material_name = 'defaultLit'
    if isinstance(mesh.visual, trimesh.visual.color.ColorVisuals):
        vertex_colors = mesh.visual.vertex_colors
        if vertex_colors is not None:
            mesh_o3d.vertex.colors = o3d.core.Tensor(
            np.divide(vertex_colors[:, :3], 255), o3d.core.float32)
        face_colors = mesh.visual.face_colors
        if face_colors is not None:
            mesh_o3d.triangle.colors = o3d.core.Tensor(
            np.divide(face_colors[:, :3], 255), o3d.core.float32)
    if isinstance(mesh.visual, trimesh.visual.texture.TextureVisuals):
        texture_uvs = mesh.visual.uv
        mesh_o3d.vertex.texture_uvs = o3d.core.Tensor(texture_uvs, o3d.core.float32)
        if isinstance(mesh.visual.material, trimesh.visual.material.PBRMaterial):
            texture = mesh.visual.material.baseColorTexture
            if texture is not None:
                mesh_o3d.material.texture_maps['albedo'] = o3d.t.geometry.Image(
                    o3d.core.Tensor(np.asarray(texture), o3d.core.uint8))
            roughness = mesh.visual.material.roughnessFactor
            if roughness is not None:
                mesh_o3d.material.scalar_properties['roughness'] = roughness
            base_color = mesh.visual.material.baseColorFactor
            if base_color is not None:
                mesh_o3d.material.vector_properties['base_color'][:] = np.divide(
                    base_color, 255)
            metallic = mesh.visual.material.metallicFactor
            if metallic is not None:
                mesh_o3d.material.scalar_properties['metallic'] = metallic

    if scale is not None:
        mesh_o3d.vertex.positions *= o3d.core.Tensor(scale, o3d.core.float32)

    return mesh_o3d

def load_geometry_object(geometry_object: pin.GeometryObject,
                         color: Optional[np.ndarray] = None
                         ) -> Optional[o3d.t.geometry.TriangleMesh]:
    """Load a single geometry object
    """
    # Extract geometry information
    geom = geometry_object.geometry
    mesh_path = geometry_object.meshPath
    texture_path = ""
    if geometry_object.overrideMaterial:
        # Get material from URDF. The color is only used if no texture or
        # if its value is not the default because meshColor is never unset.
        if os.path.exists(geometry_object.meshTexturePath):
            texture_path = geometry_object.meshTexturePath
        if color is None and (not texture_path or any(
                geometry_object.meshColor != [0.9, 0.9, 0.9, 1.0])):
            color = geometry_object.meshColor

    # Try to load mesh from path first
    is_success = True
    if os.path.exists(mesh_path):
        mesh = load_mesh_file(mesh_path, geometry_object.meshScale)
    else:
        # Each geometry must have at least a color or a texture
        if color is None and not texture_path:
            color = np.array([0.75, 0.75, 0.75, 1.0])

        # Append a primitive geometry
        if isinstance(geom, hppfcl.Capsule):
            mesh = make_capsule(geom.radius, 2.0 * geom.halfLength)
        elif isinstance(geom, hppfcl.Cylinder):
            mesh = o3d.t.geometry.TriangleMesh.create_cylinder(
                geom.radius, 2*geom.halfLength)
        elif isinstance(geom, hppfcl.Cone):
            mesh = o3d.t.geometry.TriangleMesh.create_cylinder(
                geom.radius, 2*geom.halfLength)
        elif isinstance(geom, hppfcl.Box):
            mesh = o3d.t.geometry.TriangleMesh.create_box(*(2.0*geom.halfSide))
        elif isinstance(geom, hppfcl.Sphere):
            mesh = o3d.t.geometry.TriangleMesh.create_sphere(geom.radius)
        elif isinstance(geom, (hppfcl.Convex, hppfcl.BVHModelBase)):
            # Extract vertices and faces from geometry
            if isinstance(geom, hppfcl.Convex):
                vertices = geom.points()
                num_faces, get_faces = geom.num_polygons, geom.polygons
            else:
                vertices = geom.vertices()
                num_faces, get_faces = geom.num_tris, geom.tri_indices
            faces = np.empty((num_faces, 3), dtype=np.int32)
            for i in range(num_faces):
                tri = get_faces(i)
                for j in range(3):
                    faces[i, j] = tri[j]

            # Return immediately if there is nothing to load
            if num_faces == 0:
                return None

            # Create primitive triangle geometry
            mesh = o3d.t.geometry.TriangleMesh()
            mesh.vertex.positions = o3d.core.Tensor(vertices, o3d.core.float32)
            mesh.triangle.indices = o3d.core.Tensor(faces, o3d.core.int32)
            mesh.compute_vertex_normals()
        else:
            is_success = False
        if is_success:
            mesh.material.set_default_properties()
            mesh.material.material_name = 'defaultLit'

    # Early return if impossible to load the geometry for some reason
    if not is_success:
        warnings.warn(
            f"Unsupported geometry type for {geometry_object.name} "
            f"({type(geom)})", category=UserWarning, stacklevel=2)
        return None

    # Set material
    if texture_path:
        mesh.material.texture_maps['albedo'] = o3d.io.read_image(texture_path)
    if color is not None:
        mesh.material.vector_properties['base_color'] = np.divide(
            color, 255, dtype=np.float32)

    return mesh

# data_root_dir = resource_filename("gym_jiminy.envs", "data/bipedal_robots/atlas")
# mesh = load_mesh_file(os.path.join(data_root_dir, "meshes/utorso.dae"))
# vis.draw([mesh])

env = gym.make("gym_jiminy.envs:atlas-v0")
env.reset()
# env.render()
# env.viewer.display_collisions(True)
for model, data in (
        # (env.robot.collision_model, env.robot.collision_data),
        (env.robot.visual_model, env.robot.visual_data),
    ):
    meshes = {
        geom.name: load_geometry_object(geom) if geom.name != "ground" else None
        for geom in model.geometryObjects}
    for mesh, oMg in zip(meshes.values(), data.oMg):
        if mesh is not None:
            mesh.transform(oMg.homogeneous)
    # window = vis.Visualizer()
    # for name, mesh in meshes.items():
    #     if mesh is not None:
    #         window.add_geometry(name, mesh)
    #     window.poll_events()
    #     window.update_renderer()
    vis.draw(
        [mesh for mesh in meshes.values() if mesh is not None],
        rpc_interface=True)
duburcqa commented 1 year ago

It appears to be much slower than expected, at least to load first render meshes. Once fully loaded it is reasonably fast on both integrated and dedicated GPU though. Full CPU rendering is pretty laggy but this is expected.

duburcqa commented 1 month ago

For the record, the PR https://github.com/duburcqa/jiminy/pull/826 is bringing many optimisation regarding panda3d rendering speed, for both offscreen and onscreen rendering. So it is even more likely today than replacing it by Open3d is not an option. As of today, all OS and GPU types are supported pretty well, expect for a few features for which workarounds exist. There is also an official port to web-assembly in the making. In the end, Panda3d is probably not going anywhere soon. Finally, adding one additional backend would be very time consuming for little advantage, since notebook rendering is not a serious use case and embedded video is a robust fallback at the time being. Closing.