Teleoperation and Data Collection with AirExo
import os
import time
import json
import logging
import argparse
import threading
import numpy as np
from pynput import keyboard
from easydict import EasyDict as edict
from easyrobot.camera.api import get_rgbd_camera
from easyrobot.utils.shm import SharedMemoryManager


def to_type(s):
    if s == "float":
        return np.float64
    elif s == "int":
        return np.int64
    elif s == "uint8":
        return np.uint8
    elif s == "uint16":
        return np.uint16
    else:
        raise AttributeError("Unsupported type name.")


class Collector(object):
    """
    Collector.
    """
    def __init__(self, cfgs, subpath = None, user = 0, **kwargs):
        super(Collector, self).__init__()
        # Data path
        self.path = cfgs.data_path
        if subpath is not None:
            self.path = os.path.join(self.path, subpath)
        if os.path.exists(self.path) == False:
            os.makedirs(self.path)
        # Initialize logger
        self.cfgs = cfgs
        self.logger = logging.getLogger(self.cfgs.get('logger_name', 'Collector'))
        # Initialize cameras
        self.cameras = []
        for cam in cfgs.cameras:
            cam = get_rgbd_camera(**cam)
            cam.streaming()
            self.cameras.append(cam)
        # Initialize data collectors
        self.shm_list = cfgs.shm_data.keys()
        self.shm_managers = []
        for shm_name in self.shm_list:
            self.shm_managers.append(
                SharedMemoryManager(
                    shm_name,
                    1,
                    cfgs.shm_data[shm_name]["shape"],
                    to_type(cfgs.shm_data[shm_name]["type"])
                )
            )
        self.meta = edict()
        self.meta.cfgs = cfgs
        self.meta.user = user
        # Initialize collection
        self.is_collecting = False

    def receive(self):
        '''
        Receive the data.
        '''
        res = {'time': int(time.time() * 1000)}
        for i, shm_name in enumerate(self.shm_list):
            res[shm_name] = self.shm_managers[i].execute()
        return res

    def save(self):
        '''
        Save the collected data to the data path, return the timestamp at collection.
        '''
        res = self.receive()
        timestamp = res['time']

        np.save(os.path.join(self.path, '{}.npy'.format(timestamp)), res, allow_pickle = True)
        return timestamp

    def start(self, delay_time = 0.0):
        '''
        Start collecting data.
        
        Parameters:
        - delay_time: float, optional, default: 0.0, the delay time before collecting data.
        '''
        self.thread = threading.Thread(target = self.collecting_thread, kwargs = {'delay_time': delay_time})
        self.thread.setDaemon(True)
        self.thread.start()
        self.meta.start_time = int(time.time() * 1000)
        self.meta.timestamps = []

    def collecting_thread(self, delay_time = 0.0):
        time.sleep(delay_time)
        self.is_collecting = True
        self.logger.info('[{}] Start collecting data ...'.format(self.cfgs.name))
        while self.is_collecting:
            timestamp = self.save()
            self.meta.timestamps.append(timestamp)

    def stop(self):
        '''
        Stop collecting data.
        '''
        self.meta.stop_time = int(time.time() * 1000)
        self.is_collecting = False
        self.thread.join()
        self.logger.info('[{}] Stop collecting data.'.format(self.cfgs.name))
        for cam in self.cameras:
            cam.stop_streaming()
    
    def save_meta(self):
        '''
        Save metadata.
        '''
        with open(os.path.join(self.path, 'meta.json'), 'w') as f:
            json.dump(self.meta, f)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--cfg', '-c', 
        default = os.path.join('configs', 'flexiv_teleop.yaml'), 
        help = 'path to the configuration file', 
        type = str
    )
    parser.add_argument('--task', '-t', default = 0, help = 'task id', type = int)
    parser.add_argument('--scene', '-s', default = 0, help = 'scene id', type = int)
    parser.add_argument('--user', '-u', default = 0, help = 'user id', type = int)
    args = parser.parse_args()
    if not os.path.exists(args.cfg):
        raise AttributeError('Please provide the configuration file {}.'.format(args.cfg))
    with open(args.cfg, 'r') as f:
        cfgs = edict(json.load(f))
    if cfgs.mode == 'exoskeleton':
        subpath = os.path.join('task{}_itw'.format(args.task), 'scene{}'.format(args.scene))
    else:
        subpath = os.path.join('task{}'.format(args.task), 'scene{}'.format(args.scene))
    collector = Collector(cfgs, subpath, user = args.user)

    has_stop = False

    def _on_press(key):
        global has_stop
        try:
            if key.char == 'q':
                if not has_stop:
                    collector.stop()
                    has_stop = True
        except AttributeError:
            pass

    def _on_release(key):
        pass

    listener = keyboard.Listener(on_press = _on_press, on_release = _on_release)
    collector.start()
    listener.start()
    while not has_stop:
        pass
    listener.stop()
    collector.save_meta()

collector/collector.py
AirExo/collector
import re
from easyrobot.robot.base import RobotBase
from exoskeleton.robot.flexiv import FlexivInterface


def get_robot(**params):
    name = params.get('name', None)
    try:
        if re.fullmatch('[ -_]*flexiv[ -_]*', str.lower(name)):
            return FlexivInterface(**params)
        else:
            return RobotBase(**params)
    except Exception:
        return RobotBase(**params)

exoskeleton/robot/api.py
AirExo/collector
"""
Modified flexiv robot interface for exoskeleton tele-operation.
"""
from easyrobot.robot.flexiv import FlexivRobot


class FlexivInterface(FlexivRobot):
    """
    Modified Flexiv Robot Interface.
    """
    def __init__(
        self, 
        robot_ip_address, 
        pc_ip_address, 
        gripper = {}, 
        with_streaming = False,
        streaming_freq = 30,
        shm_name = "robot",
        joint_num = 7,
        min_joint_diff = 0.01,
        **kwargs
    ):
        super(FlexivInterface, self).__init__(
            robot_ip_address = robot_ip_address,
            pc_ip_address = pc_ip_address,
            gripper = gripper,
            with_streaming = with_streaming,
            streaming_freq = streaming_freq,
            shm_name = shm_name,
            **kwargs
        )
        self.joint_num = joint_num
        self.prev_pos = [0.0] * self.DOF
        self.min_joint_diff = min_joint_diff
    
    def action(self, action_dict, wait = False):
        '''
        Perform action in the given action dict.

        Parameters:
        - wait: whether to wait until the robot reach the target position.
        '''
        joint_pos = []
        target_vel = action_dict.get('target_joint_vel', [0.0] * self.DOF)
        target_acc = action_dict.get('target_joint_acc', [0.0] * self.DOF)
        max_vel = action_dict.get('max_joint_vel', [2.0] * self.DOF)
        max_acc = action_dict.get('max_joint_acc', [3.0] * self.DOF)
        joint_pos = action_dict["enc"][:self.joint_num]

        if any(abs(self.prev_pos[i] - joint_pos[i]) > self.min_joint_diff for i in range(self.joint_num)):
            joint_pos = joint_pos
        else:
            joint_pos = self.prev_pos
        self.prev_pos = joint_pos

        self.send_joint_pos(joint_pos, wait, target_vel, target_acc, max_vel, max_acc)
        gripper_action = action_dict["enc"][self.joint_num:]
        self.gripper_action(*gripper_action)
exoskeleton/robot/flexiv.py
AirExo/collector
import os
import time
import json
import logging
import threading
from easydict import EasyDict as edict
from easyrobot.utils.degree import *
from easyrobot.utils.logger import ColoredLogger
from easyrobot.encoder.angle import AngleEncoder
from exoskeleton.robot.api import get_robot


class SingleArmTeleOperator(object):
    '''
    Tele-operator for single robot arm.
    '''
    def __init__(self, cfgs):
        '''
        Initialization.

        Parameters:
        - cfgs: the path to the configuration file.
        '''
        super(SingleArmTeleOperator, self).__init__()
        if not os.path.exists(cfgs):
            raise AttributeError('Please provide the configuration file {}.'.format(cfgs))
        with open(cfgs, 'r') as f:
            self.cfgs = edict(json.load(f))
        logging.setLoggerClass(ColoredLogger)
        self.logger = logging.getLogger(self.cfgs.get('logger_name', 'TeleOP'))
        self.action_freq = self.cfgs.freq
        self.is_teleop = False
        if not self.cfgs.encoder:
            self.ids = [1, 2, 3, 4, 5, 6, 7, 8]
            self.encoder = None
        else:
            self.encoder = AngleEncoder(**self.cfgs.encoder)
            self.ids = self.encoder.ids
            self.encoder.streaming()
        self.num_ids = len(self.ids)
        self.robot = get_robot(**self.cfgs.robot)
        self.robot.streaming()
    
    def mapping(self, enc, x, num):
        scaling = self.cfgs.mapping[enc].scaling
        emin = self.cfgs.mapping[enc].encoder_min
        emax = self.cfgs.mapping[enc].encoder_max
        edir = self.cfgs.mapping[enc].encoder_direction
        erad = self.cfgs.mapping[enc].encoder_rad
        rmin = self.cfgs.mapping[enc].robot_min
        rmax = self.cfgs.mapping[enc].robot_max
        if erad:
            x = rad_2_deg(x)
        x = deg_clip_in_range(x, emin, emax, edir)
        if scaling:
            return deg_percentile(x, emin, emax, edir) * (rmax - rmin) + rmin
        rdir = self.cfgs.mapping[enc].robot_direction
        me, mr = self.cfgs.mapping[enc].encoder_mapping, self.cfgs.mapping[enc].robot_mapping
        rrad = self.cfgs.mapping[enc].robot_rad
        rzc = self.cfgs.mapping[enc].robot_zero_centered
        fixed = self.cfgs.mapping[enc].fixed
        x = deg_clip_in_range(mr + deg_distance(me, x, edir) * rdir, rmin, rmax, rdir)
        if fixed:
            x = self.cfgs.mapping[enc].fixed_value
        if rzc:
            x = deg_zero_centered(x, rmin, rmax, rdir)
        if rrad:
            x = deg_2_rad(x)
        return x
        
    def transform_action(self, enc_res):
        '''
        Transform the action in encoder field into robot field.
        '''
        action_dict = {}
        action = []
        for i in range(self.num_ids):
            rad = self.mapping("enc{}".format(self.ids[i]), enc_res[i], i)
            action.append(rad)
        action_dict["enc"] = action
        return action_dict

    def initialization(self):
        self.logger.info('[TeleOP-{}] Initialization ... Please remain still.'.format(self.cfgs.name))
        action_dict = {}
        action = []
        for i in range(self.num_ids):
            enc = "enc{}".format(self.ids[i])
            x = self.cfgs.mapping[enc].initial_value
            scaling = self.cfgs.mapping[enc].scaling
            if scaling:
                action.append(x)
                continue
            rmin = self.cfgs.mapping[enc].robot_min
            rmax = self.cfgs.mapping[enc].robot_max
            rdir = self.cfgs.mapping[enc].robot_direction
            rrad = self.cfgs.mapping[enc].robot_rad
            rzc = self.cfgs.mapping[enc].robot_zero_centered
            if rzc:
                x = deg_zero_centered(x, rmin, rmax, rdir)
            if rrad:
                x = deg_2_rad(x)
            action.append(x)
        action_dict["enc"] = action
        action_dict["max_joint_vel"] = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
        action_dict["max_joint_acc"] = [0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]
        self.robot.action(action_dict, wait = True)
        self.logger.info('[TeleOP-{}] Finish initialization.'.format(self.cfgs.name))
    

    def calibration(self):
        '''
        Calibrate the robot with the exoskeleton when starting tele-operation.
        '''
        self.logger.info('[TeleOP-{}] Initially calibrate ... Please remain still.'.format(self.cfgs.name))
        enc_res = self.encoder.fetch_info()
        action_dict = self.transform_action(enc_res)
        action_dict["max_joint_vel"] = [0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]
        action_dict["max_joint_acc"] = [0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]
        self.robot.action(action_dict, wait = True)
        self.logger.info('[TeleOP-{}] Finish initial calibration.'.format(self.cfgs.name))
        
    def start(self, delay_time = 0.0):
        '''
        Start tele-operation.
        
        Parameters:
        - delay_time: float, optional, default: 0.0, the delay time before collecting data.
        '''
        self.thread = threading.Thread(target = self.teleop_thread, kwargs = {'delay_time': delay_time})
        self.thread.setDaemon(True)
        self.thread.start()
    
    def teleop_thread(self, delay_time = 0.0):
        time.sleep(delay_time)
        self.is_teleop = True
        self.logger.info('[{}] Start tele-operation ...'.format(self.cfgs.name))
        while self.is_teleop:
            enc_res = self.encoder.fetch_info()
            action_dict = self.transform_action(enc_res)
            self.robot.action(action_dict)
            time.sleep(1.0 / self.action_freq)
    
    def stop(self):
        '''
        Stop tele-operation process.
        '''
        self.is_teleop = False
        if self.thread:
            self.thread.join()
        self.logger.info('[{}] Stop tele-operation.'.format(self.cfgs.name))
        if self.encoder:
            self.encoder.stop_streaming()
        self.robot.stop_streaming()
        self.robot.stop()


class DualArmTeleOperator(object):
    '''
    Tele-operator for dual robot arm.
    '''
    def __init__(self, cfgs_left, cfgs_right):
        '''
        Initialization.

        Parameters:
        - cfgs_left: the path to the configuration file of the left arm;
        - cfgs_right: the path to the configuration file of the right arm.
        '''
        super(DualArmTeleOperator, self).__init__()
        self.op_left = SingleArmTeleOperator(cfgs_left)
        self.op_right = SingleArmTeleOperator(cfgs_right)
    
    def initialization(self):
        self.op_left.initialization()
        self.op_right.initialization()
        return
    
    def calibration(self):
        '''
        Calibrate the robot with the exoskeleton when starting tele-operation.
        '''
        self.op_left.calibration()
        self.op_right.calibration()
        return

    def start(self, delay_time = 0.0):
        '''
        Start tele-operation.
        
        Parameters:
        - delay_time: float, optional, default: 0.0, the delay time before collecting data.
        '''
        self.op_left.start(delay_time)
        self.op_right.start(delay_time)

    def stop(self):
        '''
        Stop the tele-operation process.
        '''
        self.op_left.stop()
        self.op_right.stop()


if __name__ == '__main__':
    op = SingleArmTeleOperator('configs/flexiv_left_gather_balls.json')
    op.calibration()
    op.start()
    time.sleep(10)
    op.stop()
    

exoskeleton/teleop.py
AirExo/collector
import os
import time
import json
import argparse
from pynput import keyboard
from easydict import EasyDict as edict
from easyrobot.encoder.angle import AngleEncoder
from exoskeleton.teleop import DualArmTeleOperator


if __name__ == '__main__':
    os.system("kill -9 `ps -ef | grep collector | grep -v grep | awk '{print $2}'`")
    os.system('rm -f /dev/shm/*')
    os.system('udevadm trigger')

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--task', '-t', 
        default = 'gather_balls', 
        help = 'task name', 
        type = str,
        choices = ['gather_balls', 'grasp_from_the_curtained_shelf']
    )
    parser.add_argument(
        '--type',
        default = 'teleop', 
        help = 'type of demonstration collection', 
        type = str,
        choices = ['teleop', 'exoskeleton']
    )
    args = parser.parse_args()

    run_path = 'configs/flexiv_' + str(args.type) + '_' + str(args.task) + '.json'

    if not os.path.exists(run_path):
        raise AttributeError('Please provide the configuration file {}.'.format(run_path))
    with open(run_path, 'r') as f:
        cfgs = edict(json.load(f))
    
    tid = int(input('Task ID: '))
    sid = int(input('Scene ID: '))
    uid = int(input('User ID: '))
    
    if cfgs.mode == 'teleop':
        op = DualArmTeleOperator(cfgs.config.left, cfgs.config.right)
        op.initialization()
        time.sleep(2)
        op.calibration()

        has_start = False
        has_stop = False
        has_record = False

        def _on_press(key):
            global has_start
            global has_stop
            global has_record
            try:
                if key.char == 's':
                    if not has_start:
                        op.start()
                        has_start = True
                    else:
                        pass  
                if key.char == 'q':
                    if not has_start:
                        pass
                    else:
                        op.stop()
                        op.initialization()
                        time.sleep(2)
                        has_stop = True
                if key.char == 'r':
                    if not has_record:
                        os.system('bash collector/collector.sh {} {} {} {} &'.format(run_path, tid, sid, uid))
                        time.sleep(1)
                        has_record = True
                    else:
                        pass     
            except Exception as e:
                print(e)
                pass

        def _on_release(key):
            pass

        listener = keyboard.Listener(on_press = _on_press, on_release = _on_release)
        listener.start()
        while not has_start:
            time.sleep(0.5)
        while not has_stop:
            time.sleep(0.5)
        listener.stop()
    elif cfgs.mode == 'exoskeleton':
        encoder_left = AngleEncoder(**cfgs.encoder_left)
        encoder_left.streaming()
        encoder_right = AngleEncoder(**cfgs.encoder_right)
        encoder_right.streaming()
        has_start = False
        has_stop = False

        def _on_press(key):
            global has_start
            global has_stop
            try:
                if key.char == 'q':
                    if not has_start:
                        pass
                    else:
                        encoder_right.stop_streaming()
                        encoder_left.stop_streaming()
                        has_stop = True
                if key.char == 's':
                    if not has_start:
                        os.system('bash collector/collector.sh {} {} {} {} &'.format(run_path, tid, sid, uid))
                        time.sleep(3)
                        has_start = True
                    else:
                        pass     
            except AttributeError:
                pass

        def _on_release(key):
            pass

        listener = keyboard.Listener(on_press = _on_press, on_release = _on_release)
        listener.start()
        while not has_stop:
            pass
        listener.stop()
        
    else:
        raise AttributeError('Invalid type.')
main.py
AirExo/collector
import numpy as np

from easyrobot.encoder.angle import AngleEncoder

encoder = AngleEncoder(ids = [1, 2, 3, 4, 5, 6, 7, 8], port = '/dev/ttyUSB0', baudrate = 115200)

while True:
    e = encoder.fetch_info()
    print(e)
    
test_encoder.py
AirExo/collector
import time
import logging
import argparse
from exoskeleton.teleop import DualArmTeleOperator


parser = argparse.ArgumentParser()
parser.add_argument(
    '--task', '-t', 
    default = 'gather_balls', 
    help = 'task name', 
    type = str,
    choices = ['gather_balls', 'grasp_from_the_curtained_shelf']
)
args = parser.parse_args()

op = DualArmTeleOperator(
    'exoskeleton/configs/flexiv_left_' + str(args.task) + '.json', 'exoskeleton/configs/flexiv_right_' + str(args.task) + '.json'
)

logger = logging.getLogger("TeleOP-left")
logger.setLevel(logging.INFO)

logger = logging.getLogger("TeleOP-right")
logger.setLevel(logging.INFO)
 
op.initialization()
time.sleep(10)

op.calibration()

time.sleep(5)

try:
    op.start()
    while True:
        time.sleep(1)
except Exception:
    op.stop()

test_exoskeleton.py
AirExo/collector
import time
import logging
import argparse
from exoskeleton.teleop import SingleArmTeleOperator


parser = argparse.ArgumentParser()
parser.add_argument(
    '--task', '-t', 
    default = 'gather_balls', 
    help = 'task name', 
    type = str,
    choices = ['gather_balls', 'grasp_from_the_curtained_shelf']
)
args = parser.parse_args()

op = SingleArmTeleOperator('exoskeleton/configs/flexiv_left_' + str(args.task) + '.json')

logger = logging.getLogger("TeleOP-left")
logger.setLevel(logging.INFO)

op.initialization()
time.sleep(5)

op.calibration()

time.sleep(5)

try:
    op.start()
    while True:
        time.sleep(1)
except Exception:
    op.stop()

test_exoskeleton_left.py
AirExo/collector
import time
import logging
import argparse
from exoskeleton.teleop import SingleArmTeleOperator


parser = argparse.ArgumentParser()
parser.add_argument(
    '--task', '-t', 
    default = 'gather_balls', 
    help = 'task name', 
    type = str,
    choices = ['gather_balls', 'grasp_from_the_curtained_shelf']
)
args = parser.parse_args()

op = SingleArmTeleOperator('exoskeleton/configs/flexiv_right_' + str(args.task) + '.json')

logger = logging.getLogger("TeleOP-right")
logger.setLevel(logging.INFO)

op.initialization()
time.sleep(5)

op.calibration()

time.sleep(5)

try:
    op.start()
    while True:
        time.sleep(1)
except Exception:
    op.stop()

test_exoskeleton_right.py
AirExo/collector