ORM querying examples¶

from tf import transformations
import itertools
import time
from typing import Optional, List, Tuple

import numpy as np

import sqlalchemy.orm
import tf
import tqdm

import pycram.orm.base
import pycram.task
from pycram.bullet_world import BulletWorld, Object as BulletWorldObject
from pycram.designators.action_designator import MoveTorsoAction, PickUpAction, NavigateAction, ParkArmsAction
from pycram.designators.object_designator import ObjectDesignatorDescription
import pycram.enums
from pycram.plan_failures import PlanFailure
from pycram.process_module import ProcessModule

from pycram.process_module import simulated_robot
import sqlalchemy.orm
# from pycram.resolver.location.jpt_location import JPTCostmapLocation
import pycram.orm
from pycram.orm.base import Position, RobotState
from pycram.orm.task import TaskTreeNode, Code
from pycram.orm.action_designator import PickUpAction as ORMPickUpAction
from pycram.orm.object_designator import Object
import sqlalchemy.sql
import pandas as pd

from pycram.pose import Pose

np.random.seed(420)

ProcessModule.execution_delay = False
pycram.orm.base.ProcessMetaData().description = "Tutorial for learning from experience in a Grasping action."


class GraspingExplorer:
    """Class to try randomized grasping plans."""

    world: Optional[BulletWorld]

    def __init__(self, robots: Optional[List[Tuple[str, str]]] = None, objects: Optional[List[Tuple[str, str]]] = None,
                 arms: Optional[List[str]] = None, grasps: Optional[List[str]] = None,
                 samples_per_scenario: int = 1000):
        """
        Create a GraspingExplorer.
        :param robots: The robots to use
        :param objects: The objects to try to grasp
        :param arms: The arms of the robot to use
        :param grasps: The grasp orientations to use
        :param samples_per_scenario: The number of tries per scenario.
        """
        # store exploration space
        if not robots:
            self.robots: List[Tuple[str, str]] = [("pr2", "pr2.urdf")]

        if not objects:
            self.objects: List[Tuple[str, pycram.enums.ObjectType, str]] = [("cereal", pycram.enums.ObjectType.BREAKFAST_CEREAL, "breakfast_cereal.stl"),
                                                                            ("bowl", pycram.enums.ObjectType.BOWL, "bowl.stl"),
                                                                            ("milk", pycram.enums.ObjectType.MILK, "milk.stl"),
                                                                            ("spoon", pycram.enums.ObjectType.SPOON, "spoon.stl")]

        if not arms:
            self.arms: List[str] = ["left", "right"]

        if not grasps:
            self.grasps: List[str] = ["left", "right", "front", "top"]

        # store trials per scenario
        self.samples_per_scenario: int = samples_per_scenario

        # chain hyperparameters
        self.hyper_parameters = [self.robots, self.objects, self.arms, self.grasps]

        self.total_tries = 0
        self.total_failures = 0

    def perform(self, session: sqlalchemy.orm.Session):
        """
        Perform all experiments.
        :param session: The database-session to insert the samples in.
        """

        # create progress bar
        progress_bar = tqdm.tqdm(
            total=np.prod([len(p) for p in self.hyper_parameters]) * self.samples_per_scenario)

        self.world = BulletWorld("DIRECT")

        # for every robot
        for robot, robot_urdf in self.robots:

            # spawn it
            robot = BulletWorldObject(robot, pycram.enums.ObjectType.ROBOT, robot_urdf)

            # for every obj
            for obj, obj_type, obj_stl in self.objects:

                # spawn it
                bw_object = BulletWorldObject(obj, obj_type, obj_stl, Pose([0, 0, 0.75], [0, 0, 0, 1]))

                # create object designator
                object_designator = ObjectDesignatorDescription(names=[obj])

                # for every arm and grasp pose
                for arm, grasp in itertools.product(self.arms, self.grasps):
                    # sample positions in 2D
                    positions = np.random.uniform([-2, -2], [2, 2], (self.samples_per_scenario, 2))

                    # for every position
                    for position in positions:

                        # set z axis to 0
                        position = [*position, 0]

                        # calculate orientation for robot to face the object
                        angle = np.arctan2(position[1], position[0]) + np.pi
                        orientation = list(transformations.quaternion_from_euler(0, 0, angle, axes="sxyz"))

                        # try to execute a grasping plan
                        with simulated_robot:

                            ParkArmsAction.Action(pycram.enums.Arms.BOTH).perform()
                            # navigate to sampled position
                            NavigateAction([Pose(position, orientation)]).resolve().perform()

                            # move torso
                            height = np.random.uniform(0., 0.33, 1)[0]
                            MoveTorsoAction.Action(height).perform()

                            # try to pick it up
                            try:
                                PickUpAction(object_designator, [arm], [grasp]).resolve().perform()

                            # if it fails
                            except PlanFailure:

                                # update failure stats
                                self.total_failures += 1

                            # reset BulletWorld
                            self.world.reset_bullet_world()

                            # update progress bar
                            self.total_tries += 1

                            # insert into database
                            pycram.task.task_tree.insert(session, use_progress_bar=False)
                            pycram.task.reset_tree()

                            progress_bar.update()
                            progress_bar.set_postfix(success_rate=(self.total_tries - self.total_failures) /
                                                                  self.total_tries)

                bw_object.remove()
            robot.remove()

engine = sqlalchemy.create_engine("sqlite+pysqlite:///:memory:")
session = sqlalchemy.orm.Session(bind=engine)
pycram.orm.base.Base.metadata.create_all(bind=engine)
session.commit()

explorer = GraspingExplorer(samples_per_scenario=30)
explorer.perform(session)

from pycram.enums import ObjectType

milk = BulletWorldObject("Milk", ObjectType.MILK, "milk.stl")

# query all relative robot positions in regard to an objects position
# make sure to order the joins() correctly
query = (session.query(ORMPickUpAction.arm, ORMPickUpAction.grasp, RobotState.torso_height, Position.x, Position.y)
         .join(TaskTreeNode.code)
         .join(Code.designator.of_type(ORMPickUpAction))
         .join(ORMPickUpAction.robot_state)
         .join(RobotState.pose)
         .join(pycram.orm.base.Pose.position)
         .join(ORMPickUpAction.object).filter(Object.type == milk.type)
                                      .filter(TaskTreeNode.status == "SUCCEEDED"))
print(query)

df = pd.read_sql_query(query.statement, session.get_bind())
print(df)

from pycram.orm.base import Pose as ORMPose

robot_pose = sqlalchemy.orm.aliased(ORMPose)
object_pose = sqlalchemy.orm.aliased(ORMPose)
robot_position = sqlalchemy.orm.aliased(Position)
object_position = sqlalchemy.orm.aliased(Position)

query = (session.query(TaskTreeNode.status, Object.type, 
                       sqlalchemy.label("relative torso height", object_position.z - RobotState.torso_height),
                       sqlalchemy.label("x", robot_position.x - object_position.x),
                       sqlalchemy.label("y", robot_position.y - object_position.y))
         .join(TaskTreeNode.code)
         .join(Code.designator.of_type(ORMPickUpAction))
         .join(ORMPickUpAction.robot_state)
         .join(robot_pose, RobotState.pose)
         .join(robot_position, robot_pose.position)
         .join(ORMPickUpAction.object)
         .join(object_pose, Object.pose)
         .join(object_position, object_pose.position))
print(query)

df = pd.read_sql(query.statement, session.get_bind())
df["status"] = df["status"].apply(lambda x: str(x.name))
print(df)