Source code for grutopia.core.robot.robot
from functools import wraps
from typing import Dict, Tuple
import numpy as np
from omni.isaac.core.prims import RigidPrim
from omni.isaac.core.robots.robot import Robot as IsaacRobot
from omni.isaac.core.scenes import Scene
from grutopia.core.config import RobotUserConfig, TaskUserConfig
from grutopia.core.robot.robot_model import RobotModel, RobotModels
from grutopia.core.util import log
[docs]class BaseRobot:
"""Base class of robot."""
robots = {}
def __init__(self, config: RobotUserConfig, robot_model: RobotModel, scene: Scene):
self.name = config.name
self.robot_model = robot_model
self.user_config = config
self.isaac_robot: IsaacRobot | None = None
self.controllers = {}
self.sensors = {}
[docs] def set_up_to_scene(self, scene: Scene):
"""Set up robot in the scene.
Args:
scene (Scene): scene to setup.
"""
config = self.user_config
robot_model = self.robot_model
scene.add(self.isaac_robot)
log.debug('self.isaac_robot: ' + str(self.isaac_robot))
from grutopia.core.robot.controller import BaseController, create_controllers
from grutopia.core.robot.sensor import BaseSensor, create_sensors
self.controllers: Dict[str, BaseController] = create_controllers(config, robot_model, self, scene)
self.sensors: Dict[str, BaseSensor] = create_sensors(config, robot_model, self, scene)
[docs] def apply_action(self, action: dict):
"""Apply actions of controllers to robot.
Args:
action (dict): action dict.
key: controller name.
value: corresponding action array.
"""
raise NotImplementedError()
[docs] def get_obs(self) -> dict:
"""Get observation of robot, including controllers, sensors, and world pose.
Raises:
NotImplementedError: _description_
"""
raise NotImplementedError()
[docs] def get_robot_ik_base(self) -> RigidPrim:
"""Get base link of ik controlled parts.
Returns:
RigidPrim: rigid prim of ik base link.
"""
raise NotImplementedError()
[docs] def get_robot_base(self) -> RigidPrim:
"""
Get base link of robot.
Returns:
RigidPrim: rigid prim of robot base link.
"""
raise NotImplementedError()
[docs] def get_robot_scale(self) -> np.ndarray:
"""Get robot scale.
Returns:
np.ndarray: robot scale in (x, y, z).
"""
return self.isaac_robot.get_local_scale()
[docs] def get_robot_articulation(self) -> IsaacRobot:
"""Get isaac robots instance (articulation).
Returns:
Robot: robot articulation.
"""
return self.isaac_robot
def get_controllers(self):
return self.controllers
def get_world_pose(self) -> Tuple[np.ndarray, np.ndarray]:
return self.isaac_robot.get_world_pose()
[docs] @classmethod
def register(cls, name: str):
"""Register a robot class with its name(decorator).
Args:
name(str): name of the robot class.
"""
def decorator(robot_class):
cls.robots[name] = robot_class
@wraps(robot_class)
def wrapped_function(*args, **kwargs):
return robot_class(*args, **kwargs)
return wrapped_function
return decorator
[docs]def create_robots(config: TaskUserConfig, robot_models: RobotModels, scene: Scene) -> Dict[str, BaseRobot]:
"""Create robot instances in config.
Args:
config (TaskUserConfig): user config.
robot_models (RobotModels): robot models.
scene (Scene): isaac scene.
Returns:
Dict[str, BaseRobot]: robot instances dictionary.
"""
robot_map = {}
for robot in config.robots:
if robot.type not in BaseRobot.robots:
raise KeyError(f'unknown robot type "{robot.type}"')
robot_cls = BaseRobot.robots[robot.type]
robot_models = robot_models.robots
r_model = None
for model in robot_models:
if model.type == robot.type:
r_model = model
if r_model is None:
raise KeyError(f'robot model of "{robot.type}" is not found')
robot_ins = robot_cls(robot, r_model, scene)
robot_map[robot.name] = robot_ins
robot_ins.set_up_to_scene(scene)
log.debug(f'===== {robot.name} loaded =====')
return robot_map