You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
733 lines
26 KiB
733 lines
26 KiB
import numpy as np
|
|
from robocasa.environments.locomanipulation.base import RobotPoseRandomizer
|
|
from robocasa.environments.locomanipulation.locomanip import LMSimpleEnv
|
|
from robocasa.utils.dexmg_utils import DexMGConfigHelper
|
|
from robocasa.utils.scene.configs import (
|
|
ObjectConfig,
|
|
ReferenceConfig,
|
|
SamplingConfig,
|
|
SceneScaleConfig,
|
|
)
|
|
from robocasa.utils.scene.scene import SceneObject
|
|
from robocasa.utils.scene.success_criteria import (
|
|
AllCriteria,
|
|
AnyCriteria,
|
|
IsGrasped,
|
|
IsInContact,
|
|
IsPositionInRange,
|
|
IsRobotInRange,
|
|
IsUpright,
|
|
NotCriteria,
|
|
SuccessCriteria,
|
|
)
|
|
from robocasa.utils.visuals_utls import Gradient, randomize_materials_rgba
|
|
|
|
|
|
class LMPickBottle(LMSimpleEnv, DexMGConfigHelper):
|
|
SCENE_SCALE = SceneScaleConfig(planar_scale=1.0)
|
|
|
|
TABLE_GRADIENT: Gradient = Gradient(
|
|
np.array([0.68, 0.34, 0.07, 1.0]), np.array([1.0, 1.0, 1.0, 1.0])
|
|
)
|
|
LIFT_OFFSET = 0.1
|
|
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
self.table = SceneObject(
|
|
ObjectConfig(
|
|
name="table",
|
|
mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
|
|
scale=1.0,
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.02, 0.02]),
|
|
y_range=np.array([-0.02, 0.02]),
|
|
reference_pos=np.array([0.5, 0, 0]),
|
|
rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
|
|
),
|
|
)
|
|
)
|
|
self.bottle = SceneObject(
|
|
ObjectConfig(
|
|
name="bottle",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.08, 0.04]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference_pos=np.array([0.4, 0, self.table.mj_obj.top_offset[2]]),
|
|
),
|
|
)
|
|
)
|
|
return [self.table, self.bottle]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return AllCriteria(
|
|
IsGrasped(self.bottle, "right"),
|
|
IsPositionInRange(self.bottle, 2, self.table.mj_obj.top_offset[2] + self.LIFT_OFFSET),
|
|
)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Pick up the bottle."
|
|
|
|
def get_object(self):
|
|
return dict(
|
|
bottle=dict(obj_name=self.bottle.mj_obj.root_body, obj_type="body"),
|
|
)
|
|
|
|
def get_subtask_term_signals(self):
|
|
signals = dict()
|
|
signals["grasp_bottle"] = int(
|
|
self._check_grasp(self.robots[0].gripper["right"], self.bottle.mj_obj.contact_geoms)
|
|
)
|
|
return signals
|
|
|
|
@staticmethod
|
|
def task_config():
|
|
task = DexMGConfigHelper.AttrDict()
|
|
task.task_spec_0.subtask_1 = dict(
|
|
object_ref="bottle",
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
task.task_spec_1.subtask_1 = dict(
|
|
object_ref=None,
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
return task.to_dict()
|
|
|
|
def _reset_internal(self):
|
|
super()._reset_internal()
|
|
|
|
if not self.deterministic_reset:
|
|
self._randomize_table_rgba()
|
|
|
|
def _randomize_table_rgba(self):
|
|
randomize_materials_rgba(
|
|
rng=self.rng, mjcf_obj=self.table.mj_obj, gradient=self.TABLE_GRADIENT, linear=True
|
|
)
|
|
|
|
|
|
class LMPickBottleHigh(LMPickBottle):
|
|
TABLE_OFFSET = 0.1
|
|
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
self.table = SceneObject(
|
|
ObjectConfig(
|
|
name="table",
|
|
mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
|
|
scale=1.0,
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.02, 0.02]),
|
|
y_range=np.array([-0.02, 0.02]),
|
|
reference_pos=np.array([0.5, 0, self.TABLE_OFFSET]),
|
|
rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
|
|
),
|
|
)
|
|
)
|
|
self.bottle = SceneObject(
|
|
ObjectConfig(
|
|
name="bottle",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.08, 0.04]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference_pos=np.array(
|
|
[0.4, 0, self.TABLE_OFFSET + self.table.mj_obj.top_offset[2]]
|
|
),
|
|
reference=ReferenceConfig(obj=self.table),
|
|
),
|
|
)
|
|
)
|
|
return [self.table, self.bottle]
|
|
|
|
|
|
class LMNavPickBottle(LMPickBottle):
|
|
def _reset_internal(self):
|
|
super()._reset_internal()
|
|
|
|
if not self.deterministic_reset:
|
|
RobotPoseRandomizer.set_pose(self, (-0.3, -0.16), (-0.2, 0.2), (-np.pi / 6, np.pi / 6))
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Walk forward and pick up the bottle from the table."
|
|
|
|
|
|
class LMPickBottleGround(LMPickBottle):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
self.bottle = SceneObject(
|
|
ObjectConfig(
|
|
name="bottle",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.08, 0.04]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference_pos=np.array(
|
|
[0.4, 0, 0.075]
|
|
), # Base position on ground (z=0.075 is bottle radius)
|
|
),
|
|
)
|
|
)
|
|
return [self.bottle]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return AllCriteria(
|
|
IsGrasped(self.bottle, "right"),
|
|
IsPositionInRange(self.bottle, 2, self.LIFT_OFFSET, 10),
|
|
)
|
|
|
|
def _randomize_table_rgba(self):
|
|
pass
|
|
|
|
|
|
class LMPnPBottle(LMPickBottle):
|
|
LIFT_OFFSET = 0.15
|
|
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.table_target = SceneObject(
|
|
ObjectConfig(
|
|
name="table_target",
|
|
mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
|
|
scale=1.0,
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.02, 0.02]),
|
|
y_range=np.array([-0.02, 0.02]),
|
|
reference_pos=np.array([0.5, 1.2, 0]),
|
|
rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
|
|
),
|
|
)
|
|
)
|
|
return [self.table, self.table_target, self.bottle]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return AllCriteria(
|
|
IsUpright(self.bottle, symmetric=True), IsInContact(self.bottle, self.table_target)
|
|
)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Pick up the bottle and place it on the other table."
|
|
|
|
def get_object(self):
|
|
return dict(
|
|
bottle=dict(obj_name=self.bottle.mj_obj.root_body, obj_type="body"),
|
|
target_table=dict(obj_name=self.table_target.mj_obj.root_body, obj_type="body"),
|
|
)
|
|
|
|
def get_subtask_term_signals(self):
|
|
obj_z = self.sim.data.body_xpos[self.obj_body_id(self.bottle.mj_obj.name)][2]
|
|
target_table_pos = self.sim.data.body_xpos[self.obj_body_id(self.table_target.mj_obj.name)]
|
|
target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
|
|
return dict(obj_off_table=int(obj_z - target_table_z > self.LIFT_OFFSET))
|
|
|
|
@staticmethod
|
|
def task_config():
|
|
task = DexMGConfigHelper.AttrDict()
|
|
task.task_spec_0.subtask_1 = dict(
|
|
object_ref="bottle",
|
|
subtask_term_signal="obj_off_table",
|
|
subtask_term_offset_range=(5, 10),
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
# Second subtask for placing on target table
|
|
task.task_spec_0.subtask_2 = dict(
|
|
object_ref="target_table",
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
task.task_spec_1.subtask_1 = dict(
|
|
object_ref=None,
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
return task.to_dict()
|
|
|
|
def _randomize_table_rgba(self):
|
|
for table in [self.table_target, self.table]:
|
|
randomize_materials_rgba(
|
|
rng=self.rng, mjcf_obj=table.mj_obj, gradient=self.TABLE_GRADIENT, linear=True
|
|
)
|
|
|
|
|
|
class LMPickMultipleBottles(LMPickBottle):
|
|
BOTTLE_COLOURS = [(0.3, 0.7, 0.8, 1.0), (0.8, 0.4, 0.3, 1.0)]
|
|
BOTTLES_COUNT = 2
|
|
Y_OFFSET_STEP = 0.1
|
|
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
self.table = SceneObject(
|
|
ObjectConfig(
|
|
name="table",
|
|
mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
|
|
scale=1.0,
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.02, 0.02]),
|
|
y_range=np.array([-0.02, 0.02]),
|
|
reference_pos=np.array([0.5, 0, 0]),
|
|
rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
|
|
),
|
|
)
|
|
)
|
|
|
|
self.bottles = []
|
|
offsets = np.arange(self.BOTTLES_COUNT) - (self.BOTTLES_COUNT - 1) / 2.0
|
|
for i in range(self.BOTTLES_COUNT):
|
|
reference_pos = np.array([0.4, 0, self.table.mj_obj.top_offset[2]])
|
|
reference_pos += np.array([0, self.Y_OFFSET_STEP * offsets[i], 0])
|
|
bottle = SceneObject(
|
|
ObjectConfig(
|
|
name=f"bottle_{i}",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.08, 0.04]),
|
|
y_range=np.array([-0.04, 0.04]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference_pos=reference_pos,
|
|
),
|
|
rgba=self.BOTTLE_COLOURS[i % len(self.BOTTLE_COLOURS)],
|
|
)
|
|
)
|
|
self.bottles.append(bottle)
|
|
return [self.table, *self.bottles]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
criteria = []
|
|
for bottle in self.bottles:
|
|
criteria.append(AnyCriteria(IsGrasped(bottle, "right"), IsGrasped(bottle, "left")))
|
|
criteria.append(
|
|
IsPositionInRange(bottle, 2, self.table.mj_obj.top_offset[2] + self.LIFT_OFFSET, 10)
|
|
)
|
|
return AllCriteria(*criteria)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Pick up bottles."
|
|
|
|
def get_object(self):
|
|
return {
|
|
bottle.mj_obj.name: dict(obj_name=bottle.mj_obj.root_body, obj_type="body")
|
|
for bottle in self.bottles
|
|
}
|
|
|
|
def get_subtask_term_signals(self):
|
|
return {
|
|
f"grasp_{bottle.mj_obj.name}": int(
|
|
self._check_grasp(self.robots[0].gripper["right"], bottle.mj_obj)
|
|
or self._check_grasp(self.robots[0].gripper["left"], bottle.mj_obj)
|
|
)
|
|
for bottle in self.bottles
|
|
}
|
|
|
|
@staticmethod
|
|
def task_config():
|
|
task = DexMGConfigHelper.AttrDict()
|
|
for i in range(LMPickMultipleBottles.BOTTLES_COUNT):
|
|
subtask = dict(
|
|
object_ref=f"bottle_{i}",
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
setattr(task.task_spec_0, f"subtask_{i+1}", subtask)
|
|
task.task_spec_1.subtask_1 = dict(
|
|
object_ref=None,
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
return task.to_dict()
|
|
|
|
|
|
class LMPnPMultipleBottles(LMPickMultipleBottles):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.table_target = SceneObject(
|
|
ObjectConfig(
|
|
name="table_target",
|
|
mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
|
|
scale=1.0,
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.02, 0.02]),
|
|
y_range=np.array([-0.02, 0.02]),
|
|
reference_pos=np.array([0.5, 1.2, 0]),
|
|
rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
|
|
),
|
|
)
|
|
)
|
|
return [self.table, self.table_target, *self.bottles]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
criteria = [
|
|
AllCriteria(IsInContact(bottle, self.table_target), IsUpright(bottle, symmetric=True))
|
|
for bottle in self.bottles
|
|
]
|
|
return AllCriteria(*criteria)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Pick up bottles from one table and place it on the other."
|
|
|
|
@staticmethod
|
|
def task_config():
|
|
task = DexMGConfigHelper.AttrDict()
|
|
for i in range(LMPnPMultipleBottles.BOTTLES_COUNT):
|
|
bottle_name = f"bottle_{i}"
|
|
subtask = dict(
|
|
object_ref=bottle_name,
|
|
subtask_term_signal=f"{bottle_name}_off_table",
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
setattr(task.task_spec_0, f"subtask_{i+1}", subtask)
|
|
# Next subtask for placing on target table
|
|
task.task_spec_0.subtask_3 = dict(
|
|
object_ref="target_table",
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
task.task_spec_1.subtask_1 = dict(
|
|
object_ref=None,
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
return task.to_dict()
|
|
|
|
def get_subtask_term_signals(self):
|
|
signals = dict()
|
|
for bottle in self.bottles:
|
|
obj_z = self.sim.data.body_xpos[self.obj_body_id(self.bottle.mj_obj.name)][2]
|
|
target_table_pos = self.sim.data.body_xpos[
|
|
self.obj_body_id(self.table_target.mj_obj.name)
|
|
]
|
|
target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
|
|
signals[f"{bottle.mj_obj.name}_off_table"] = int(
|
|
obj_z - target_table_z > self.LIFT_OFFSET
|
|
)
|
|
return signals
|
|
|
|
def _randomize_table_rgba(self):
|
|
for table in [self.table_target, self.table]:
|
|
randomize_materials_rgba(
|
|
rng=self.rng, mjcf_obj=table.mj_obj, gradient=self.TABLE_GRADIENT, linear=True
|
|
)
|
|
|
|
|
|
class LMPickBottleShelf(LMPickBottle):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.shelf = SceneObject(
|
|
ObjectConfig(
|
|
name="shelf",
|
|
mjcf_path="objects/omniverse/locomanip/lab_shelf/model.xml",
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
rotation=np.array([np.pi / 2, np.pi / 2]),
|
|
reference_pos=np.array([0.9, 0, 0]),
|
|
),
|
|
)
|
|
)
|
|
self.bottle = SceneObject(
|
|
ObjectConfig(
|
|
name="bottle",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.14, -0.06]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference=ReferenceConfig(self.shelf, spawn_id=2),
|
|
),
|
|
)
|
|
)
|
|
return [self.shelf, self.bottle]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return AllCriteria(
|
|
IsGrasped(self.bottle, "right"),
|
|
NotCriteria(IsInContact(self.bottle, self.shelf)),
|
|
)
|
|
|
|
|
|
class LMNavPickBottleShelf(LMPickBottleShelf):
|
|
ROBOT_DISTANCE_THRESHOLD = 1.0
|
|
|
|
def _reset_internal(self):
|
|
super()._reset_internal()
|
|
if not self.deterministic_reset:
|
|
RobotPoseRandomizer.set_pose(self, (-0.1, 0.1), (-0.1, 0.1), (-np.pi / 6, np.pi / 6))
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return AllCriteria(
|
|
NotCriteria(IsRobotInRange(self.shelf, self.ROBOT_DISTANCE_THRESHOLD, True)),
|
|
IsGrasped(self.bottle, "right"),
|
|
NotCriteria(IsInContact(self.bottle, self.shelf)),
|
|
)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Pick up the bottle from the shelf and move backward away from it."
|
|
|
|
|
|
class LMPickBottleShelfLow(LMPickBottleShelf):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.bottle = SceneObject(
|
|
ObjectConfig(
|
|
name="bottle",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.14, -0.06]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference=ReferenceConfig(self.shelf, spawn_id=1),
|
|
),
|
|
)
|
|
)
|
|
return [self.shelf, self.bottle]
|
|
|
|
|
|
class LMNavPickBottleShelfLow(LMNavPickBottleShelf):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.bottle = SceneObject(
|
|
ObjectConfig(
|
|
name="bottle",
|
|
mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
|
|
static=False,
|
|
scale=0.6,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.14, -0.06]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference=ReferenceConfig(self.shelf, spawn_id=1),
|
|
),
|
|
)
|
|
)
|
|
return [self.shelf, self.bottle]
|
|
|
|
|
|
class LMPnPBottleToPlate(LMPnPBottle):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.plate = SceneObject(
|
|
ObjectConfig(
|
|
name="plate",
|
|
mjcf_path="objects/omniverse/locomanip/plate_1/model.xml",
|
|
scale=1.0,
|
|
static=True,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.2 - 0.08, -0.2 + 0.04]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference=ReferenceConfig(self.table_target),
|
|
),
|
|
)
|
|
)
|
|
return [self.table, self.table_target, self.bottle, self.plate]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return AllCriteria(
|
|
IsUpright(self.bottle, symmetric=True), IsInContact(self.bottle, self.plate)
|
|
)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "Pick up the bottle and place it on the plate."
|
|
|
|
def get_object(self):
|
|
return dict(
|
|
bottle=dict(obj_name=self.bottle.mj_obj.root_body, obj_type="body"),
|
|
plate=dict(obj_name=self.plate.mj_obj.root_body, obj_type="body"),
|
|
)
|
|
|
|
def get_subtask_term_signals(self):
|
|
obj_z = self.sim.data.body_xpos[self.obj_body_id(self.bottle.mj_obj.name)][2]
|
|
target_table_pos = self.sim.data.body_xpos[self.obj_body_id(self.table_target.mj_obj.name)]
|
|
target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
|
|
return dict(obj_off_table=int(obj_z - target_table_z > self.LIFT_OFFSET))
|
|
|
|
@staticmethod
|
|
def task_config():
|
|
task = DexMGConfigHelper.AttrDict()
|
|
task.task_spec_0.subtask_1 = dict(
|
|
object_ref="bottle",
|
|
subtask_term_signal="obj_off_table",
|
|
subtask_term_offset_range=(5, 10),
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
# Second subtask for placing on plate
|
|
task.task_spec_0.subtask_2 = dict(
|
|
object_ref="plate",
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
task.task_spec_1.subtask_1 = dict(
|
|
object_ref=None,
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
return task.to_dict()
|
|
|
|
|
|
class LMPnPAppleToPlate(LMPnPBottleToPlate):
|
|
def _get_objects(self) -> list[SceneObject]:
|
|
super()._get_objects()
|
|
self.apple = SceneObject(
|
|
ObjectConfig(
|
|
name="apple",
|
|
mjcf_path="objects/omniverse/locomanip/apple_0/model.xml",
|
|
static=False,
|
|
scale=1.0,
|
|
sampler_config=SamplingConfig(
|
|
x_range=np.array([-0.08, 0.04]),
|
|
y_range=np.array([-0.08, 0.08]),
|
|
rotation=np.array([-np.pi, np.pi]),
|
|
reference_pos=np.array([0.4, 0, self.table.mj_obj.top_offset[2]]),
|
|
),
|
|
rgba=(0.85, 0.1, 0.1, 1.0),
|
|
)
|
|
)
|
|
return [self.table, self.table_target, self.apple, self.plate]
|
|
|
|
def _get_success_criteria(self) -> SuccessCriteria:
|
|
return IsInContact(self.apple, self.plate)
|
|
|
|
def _get_instruction(self) -> str:
|
|
return "pick up the apple, walk left and place the apple on the plate."
|
|
|
|
def get_object(self):
|
|
return dict(
|
|
apple=dict(obj_name=self.apple.mj_obj.root_body, obj_type="body"),
|
|
plate=dict(obj_name=self.plate.mj_obj.root_body, obj_type="body"),
|
|
)
|
|
|
|
def get_subtask_term_signals(self):
|
|
obj_z = self.sim.data.body_xpos[self.obj_body_id(self.apple.mj_obj.name)][2]
|
|
target_table_pos = self.sim.data.body_xpos[self.obj_body_id(self.table_target.mj_obj.name)]
|
|
target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
|
|
return dict(obj_off_table=int(obj_z - target_table_z > self.LIFT_OFFSET))
|
|
|
|
@staticmethod
|
|
def task_config():
|
|
task = DexMGConfigHelper.AttrDict()
|
|
task.task_spec_0.subtask_1 = dict(
|
|
object_ref="apple",
|
|
subtask_term_signal="obj_off_table",
|
|
subtask_term_offset_range=(5, 10),
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
# Second subtask for placing on plate
|
|
task.task_spec_0.subtask_2 = dict(
|
|
object_ref="plate",
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
task.task_spec_1.subtask_1 = dict(
|
|
object_ref=None,
|
|
subtask_term_signal=None,
|
|
subtask_term_offset_range=None,
|
|
selection_strategy="random",
|
|
selection_strategy_kwargs=None,
|
|
action_noise=0.05,
|
|
num_interpolation_steps=5,
|
|
num_fixed_steps=0,
|
|
apply_noise_during_interpolation=False,
|
|
)
|
|
return task.to_dict()
|