Source code for concepts.benchmark.manip_tabletop.bimanual_abc.envs.bimanual_abc_lift

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : bimanual_abc_lift.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/15/2024
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Optional
from concepts.math.rotationlib_xyzw import rpy, quat_mul
from concepts.simulator.pybullet.client import BulletClient
from concepts.benchmark.manip_tabletop.bimanual_abc.bimanual_abc_env_base import BimanualABCEnvBase

__all__ = [
    'BimanualABCLiftObjectEnvBase', 'BimanualABCLiftBoxEnv',
    'BimanualABCLiftTEnv',
    'BimanualABCLiftArialTEnv', 'BimanualABCLiftArialAnythingEnv'
]


[docs] class BimanualABCLiftObjectEnvBase(BimanualABCEnvBase): def _get_lifting_bar_length(self) -> float: return 0.15 def _reset_objects(self, metainfo: dict): lifting_bar_length = self._get_lifting_bar_length() box_id1 = self.add_box((0.025, lifting_bar_length, 0.025), location_2d=(0.35, -.2), color_rgba=(0, 1, 0, 1), name='box1') box_id2 = self.add_box((0.025, lifting_bar_length, 0.025), location_2d=(0.65, -.2), color_rgba=(0, 0, 1, 1), name='box2') metainfo['objects']['box1'] = {'id': box_id1} metainfo['objects']['box2'] = {'id': box_id2}
[docs] class BimanualABCLiftBoxEnv(BimanualABCLiftObjectEnvBase): def _reset_objects(self, metainfo: dict): super()._reset_objects(metainfo) box_id = self.add_box((0.1, 0.2, 0.1), location_2d=(0.5, 0.0)) metainfo['objects']['box'] = {'id': box_id}
[docs] class BimanualABCLiftTEnv(BimanualABCLiftObjectEnvBase): def _reset_objects(self, metainfo: dict): super()._reset_objects(metainfo) T_id = self.add_t_shape( size1_2d=(0.2, 0.1), size2_2d=(0.05, 0.4), thickness=0.05, pos=(0.5, 0.0, 0.1), # The center of the shape is at the center of the handle. quat=rpy(0, -90, 0), ) metainfo['objects']['T'] = {'id': T_id}
[docs] class BimanualABCLiftArialTEnv(BimanualABCLiftObjectEnvBase): def _get_lifting_bar_length(self) -> float: return 0.25 def _reset_objects(self, metainfo: dict): super()._reset_objects(metainfo) default_quat = quat_mul(rpy(0, 0, 90), rpy(0, 180, 0)) T_id = self.add_alphabet_arial_object('T', location_3d=(0.5, 0.0, 0.1), quat=quat_mul(default_quat, rpy(0, -45, 0)), scale=1.5) metainfo['objects']['T'] = {'id': T_id}
[docs] class BimanualABCLiftArialAnythingEnv(BimanualABCLiftObjectEnvBase):
[docs] def __init__(self, alphabet, client: Optional[BulletClient] = None, is_gui: bool = True, seed: int = 1234): super().__init__(client, is_gui, seed) self.alphabet = alphabet assert len(self.alphabet) == 1 and self.alphabet.isalpha()
def _get_lifting_bar_length(self) -> float: return 0.25 def _reset_objects(self, metainfo: dict): super()._reset_objects(metainfo) default_quat = quat_mul(rpy(0, 0, 90), rpy(0, 180, 0)) letter_id = self.add_alphabet_arial_object(self.alphabet, location_3d=(0.5, 0.0, 0.1), quat=quat_mul(default_quat, rpy(0, -45, 0)), scale=1.5) metainfo['objects']['letter'] = {'id': letter_id}