Skip to content

Commit

Permalink
merge in refactor and updates for utils
Browse files Browse the repository at this point in the history
  • Loading branch information
StoneT2000 committed Dec 13, 2023
1 parent 4daf06a commit 1338e8c
Show file tree
Hide file tree
Showing 28 changed files with 13,246 additions and 3 deletions.
12 changes: 12 additions & 0 deletions mani_skill2/utils/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# ManiSkill Utilities

Description of the main modules are as follows:

`building/` - All useful utility code for building a task and/or scene. Includes functions for loading assets and articulations from various datasets, and randomization functions useful for randomizing task initialization.
`scene_builder/` - Contains code relating to the `SceneBuilder` class and provides some prebuilt scene builders for a standard table top scene, a fully articulated hand-made kitchen scene, as well as scenes from AI2THOR.

`geometry/` - Various functions for working with geometry, from sampling primitive shapes to getting axis-aligned bounding boxes of articulations/actors.

`wrappers/` - Wrapper classes that provide additional functionality such as recording videos/episodes, modifying observation spaces, as well as adapting the environment API so that RL libraries such as [Stable Baselines 3](https://github.com/DLR-RM/stable-baselines3) work out of the box.

`visualization/` - Visualization tools
275 changes: 275 additions & 0 deletions mani_skill2/utils/building/actors.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,275 @@
"""
Useful utilities for adding any object and geometry into a scene
"""

import os.path as osp
from pathlib import Path
from typing import Dict

import numpy as np
import sapien
import sapien.physx as physx
import sapien.render
import transforms3d

from mani_skill2 import ASSET_DIR
from mani_skill2.utils.io_utils import load_json

# map model dataset to a database of models
model_dbs: Dict[str, Dict[str, Dict]] = {}


# Primitive Shapes
def build_cube(
scene: sapien.Scene,
half_size: float,
color,
name: str,
dynamic: bool = True,
add_collision: bool = True,
):
builder = scene.create_actor_builder()
if add_collision:
builder.add_box_collision(
half_size=[half_size] * 3,
)
builder.add_box_visual(
half_size=[half_size] * 3,
material=sapien.render.RenderMaterial(
base_color=color,
),
)
if dynamic:
actor = builder.build(name=name)
else:
actor = builder.build_static(name=name)
return actor


def build_box(
scene: sapien.Scene,
half_sizes,
color,
name: str,
dynamic: bool = True,
add_collision: bool = True,
):
builder = scene.create_actor_builder()
if add_collision:
builder.add_box_collision(
half_size=half_sizes,
)
builder.add_box_visual(
half_size=half_sizes,
material=sapien.render.RenderMaterial(
base_color=color,
),
)
if dynamic:
actor = builder.build(name=name)
else:
actor = builder.build_static(name=name)
return actor


def build_sphere(
scene: sapien.Scene,
radius: float,
color,
name: str,
dynamic: bool = True,
add_collision: bool = True,
):
builder = scene.create_actor_builder()
if add_collision:
builder.add_sphere_collision(
radius=radius,
)
builder.add_sphere_visual(
radius=radius,
material=sapien.render.RenderMaterial(
base_color=color,
),
)
if dynamic:
actor = builder.build(name=name)
else:
actor = builder.build_static(name=name)
return actor


def build_twocolor_peg(
scene: sapien.Scene,
length,
width,
color_1,
color_2,
name: str,
dynamic: bool = True,
add_collision: bool = True,
):
builder = scene.create_actor_builder()
if add_collision:
builder.add_box_collision(
half_size=[length, width, width],
)
builder.add_box_visual(
pose=sapien.Pose(p=[-length / 2, 0, 0]),
half_size=[length / 2, width, width],
material=sapien.render.RenderMaterial(
base_color=color_1,
),
)
builder.add_box_visual(
pose=sapien.Pose(p=[length / 2, 0, 0]),
half_size=[length / 2, width, width],
material=sapien.render.RenderMaterial(
base_color=color_2,
),
)
if dynamic:
actor = builder.build(name=name)
else:
actor = builder.build_static(name=name)
return actor


RED_COLOR = [220 / 255, 12 / 255, 12 / 255, 1]
BLUE_COLOR = [0 / 255, 44 / 255, 193 / 255, 1]
GREEN_COLOR = [17 / 255, 190 / 255, 70 / 255, 1]


def build_fourcolor_peg(
scene: sapien.Scene,
length,
width,
name: str,
color_1=RED_COLOR,
color_2=BLUE_COLOR,
color_3=GREEN_COLOR,
color_4=[1, 1, 1, 1],
dynamic: bool = True,
add_collision: bool = True,
):
builder = scene.create_actor_builder()
if add_collision:
builder.add_box_collision(
half_size=[length, width, width],
)
builder.add_box_visual(
pose=sapien.Pose(p=[-length / 2, -width / 2, 0]),
half_size=[length / 2, width / 2, width],
material=sapien.render.RenderMaterial(
base_color=color_1,
),
)
builder.add_box_visual(
pose=sapien.Pose(p=[length / 2, -width / 2, 0]),
half_size=[length / 2, width / 2, width],
material=sapien.render.RenderMaterial(
base_color=color_2,
),
)
builder.add_box_visual(
pose=sapien.Pose(p=[-length / 2, width / 2, 0]),
half_size=[length / 2, width / 2, width],
material=sapien.render.RenderMaterial(
base_color=color_3,
),
)
builder.add_box_visual(
pose=sapien.Pose(p=[length / 2, width / 2, 0]),
half_size=[length / 2, width / 2, width],
material=sapien.render.RenderMaterial(
base_color=color_4,
),
)
if dynamic:
actor = builder.build(name=name)
else:
actor = builder.build_static(name=name)
return actor


### Load individual assets ###


def build_actor(model_id: str, scene: sapien.Scene, name: str):
# TODO (stao): parse model id and determine from which dataset to pull asset from
# e.g. YCB or our own sapien asset database in the future
pass


def build_actor_ycb(
model_id: str,
scene: sapien.Scene,
name: str,
root_dir=ASSET_DIR / "mani_skill2_ycb",
):
if "YCB" not in model_dbs:
_load_ycb_dataset()
model_db = model_dbs["YCB"]["model_data"] # TODO (stao): remove hardcode

builder = scene.create_actor_builder()

density = model_db[model_id].get("density", 1000)
model_scales = model_db[model_id].get("scales", [1.0])
scale = model_scales[0]
physical_material = None

model_dir = Path(root_dir) / "models" / model_id
collision_file = str(model_dir / "collision.obj")
builder.add_multiple_convex_collisions_from_file(
filename=collision_file,
scale=[scale] * 3,
material=physical_material,
density=density,
decomposition="coacd",
)

visual_file = str(model_dir / "textured.obj")
builder.add_visual_from_file(filename=visual_file, scale=[scale] * 3)
actor = builder.build(name=name)
return actor


def _load_ycb_dataset():
# load YCB if used
model_dbs["YCB"] = {
"model_data": load_json(ASSET_DIR / "mani_skill2_ycb/info_pick_v0.json"),
"builder": build_actor_ycb,
}


def build_actor_ai2(
model_id: str,
scene: sapien.Scene,
name: str,
kinematic: bool = False,
):
model_path = (
Path(ASSET_DIR)
/ "scene_datasets/ai2thor/ai2thorhab-uncompressed/assets/objects"
/ f"{model_id}.glb"
)
actor_id = name
builder = scene.create_actor_builder()
builder.add_visual_from_file(str(model_path))
if kinematic:
builder.add_nonconvex_collision_from_file(str(model_path))
actor = builder.build_kinematic(name=actor_id)
else:
builder.add_multiple_convex_collisions_from_file(
str(model_path), decomposition="coacd"
)
actor = builder.build(name=actor_id)

q = transforms3d.quaternions.axangle2quat(np.array([1, 0, 0]), theta=np.deg2rad(90))
pose = sapien.Pose(q=q)
actor.set_pose(pose)
aabb = actor.find_component_by_type(
sapien.render.RenderBodyComponent
).compute_global_aabb_tight()
height = aabb[1, 1] - aabb[0, 1]
actor.set_pose(sapien.Pose(p=[0, 0, 0], q=actor.pose.q))
return actor
Loading

0 comments on commit 1338e8c

Please sign in to comment.