sim.py

import os
import math
import json
import time
import fusion
import datetime
import pybullet
import numpy as np
import pybullet_data
from utils import common
from fusion import TSDFVolume
import pybullet_utils.bullet_client as bc


class ArmPybulletSim(object):
    def __init__(self, gui_enabled, num_cam):

        # robot viewable range
        self._view_bounds = np.array([[-0.45, 0.45],
                                      [-0.45, 0.45],
                                      [ 0.000, 0.84]]) # 3x2 rows: x,y,z cols: min,max. hard


        self._volume_bounds = self._view_bounds
        self._voxel_size = 0.005

        if gui_enabled:
            self.bullet_client = bc.BulletClient(connection_mode=pybullet.GUI)
        else:
            self.bullet_client = bc.BulletClient(connection_mode=pybullet.DIRECT)

        self.bullet_client.resetSimulation()
        self.bullet_client.setAdditionalSearchPath(pybullet_data.getDataPath())
        self.bullet_client.setGravity(0, 0, -9.8)

        self._plane_id = self.bullet_client.loadURDF("plane.urdf")
        self.bullet_client.changeDynamics(self._plane_id, -1, lateralFriction=1.0)

        # add RGB-D camera (mimic RealSense D435i)
        self._scene_cam_lookat = self._view_bounds.mean(axis=1)
        self._scene_cam_position = [self._scene_cam_lookat[0], self._scene_cam_lookat[1], 0.5]
        self._scene_cam_up_direction = [0, 0, 1]
        self._scene_cam_image_size = (480, 640)
        self._scene_cam_z_near = 0.01
        self._scene_cam_z_far = 10.0
        self._scene_cam_fov_w = 69.40
        self._scene_cam_focal_length = (float(self._scene_cam_image_size[1])/2)/np.tan((np.pi*self._scene_cam_fov_w/180)/2)
        self._scene_cam_fov_h = (math.atan((float(self._scene_cam_image_size[0])/2)/self._scene_cam_focal_length)*2/np.pi)*180
        self._scene_cam_projection_matrix = self.bullet_client.computeProjectionMatrixFOV(
            fov=self._scene_cam_fov_h,
            aspect=float(self._scene_cam_image_size[1])/float(self._scene_cam_image_size[0]),
            nearVal=self._scene_cam_z_near, farVal=self._scene_cam_z_far
        )
        self._scene_cam_intrinsics = np.array([[self._scene_cam_focal_length, 0, float(self._scene_cam_image_size[1])/2],
                                             [0, self._scene_cam_focal_length, float(self._scene_cam_image_size[0])/2],
                                             [0, 0, 1]])

        self._num_cam = num_cam
        self.debug_cam = False
        self._num_joints = 5
        self._n_steps = 0
        self._timeout = 1 # in seconds
        self._saved_robot_state_dict = {}
        self._ctrl_mode = pybullet.POSITION_CONTROL
        self.debug_cam_sphere_urdf = "assets/widowx_arm_description/urdf/weightless_sphere.urdf"
        self._robot_urdf = "assets/interbotix_descriptions/urdf/wx200_gripper.urdf"
        self.save_mesh_folder = "saved_meshes"
        common.mkdir(self.save_mesh_folder)

    # Get latest RGB-D image from scene camera
    def get_scene_cam_data(self, cam_position=None, cam_lookat=None, cam_up_direction=None):
        if cam_position is None:
            cam_position = self._scene_cam_position
        if cam_lookat is None:
            cam_lookat = self._scene_cam_lookat
        if cam_up_direction is None:
            cam_up_direction = self._scene_cam_up_direction
        cam_view_matrix = self.bullet_client.computeViewMatrix(cam_position, cam_lookat, cam_up_direction)

        cam_pose_matrix = np.linalg.inv(np.array(cam_view_matrix).reshape(4, 4).T)
        cam_pose_matrix[:, 1:3] = -cam_pose_matrix[:, 1:3]

        camera_data = self.bullet_client.getCameraImage(
            self._scene_cam_image_size[1],
            self._scene_cam_image_size[0],
            cam_view_matrix,
            self._scene_cam_projection_matrix,
            shadow=1,
            flags=self.bullet_client.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
            renderer=self.bullet_client.ER_BULLET_HARDWARE_OPENGL
        )
        rgb_pixels = np.array(camera_data[2]).reshape((self._scene_cam_image_size[0], self._scene_cam_image_size[1], 4))
        color_image = rgb_pixels[:,:,:3] # remove alpha channel
        z_buffer = np.array(camera_data[3]).reshape((self._scene_cam_image_size[0], self._scene_cam_image_size[1]))
        segmentation_mask = np.array(camera_data[4], np.int) # - not implemented yet with renderer=p.ER_BULLET_HARDWARE_OPENGL
        depth_image = (2.0*self._scene_cam_z_near*self._scene_cam_z_far)/(self._scene_cam_z_far+self._scene_cam_z_near-(2.0*z_buffer-1.0)*(self._scene_cam_z_far-self._scene_cam_z_near))
        return color_image, depth_image, segmentation_mask, cam_pose_matrix

    def reset(self):
        # load robot
        robot_start_pos = [0, 0, 0.01]
        robot_start_ori = pybullet.getQuaternionFromEuler([0, 0, 0])
        self._robot_id = self.bullet_client.loadURDF(self._robot_urdf, robot_start_pos, robot_start_ori, useFixedBase=True, flags=self.bullet_client.URDF_USE_SELF_COLLISION)

        # return observation: robot state, scene observation
        observation = self._get_observation(mesh_idx=0, if_save=False)
        return observation
    
    def reset_everything(self):
        self.bullet_client.resetSimulation()
        self.bullet_client.setAdditionalSearchPath(pybullet_data.getDataPath())
        self.bullet_client.setGravity(0, 0, -9.8)
        
        self._plane_id = self.bullet_client.loadURDF("plane.urdf")
        self.bullet_client.changeDynamics(self._plane_id, -1, lateralFriction=1.0)

        robot_start_pos = [0, 0, 0.01]
        robot_start_ori = pybullet.getQuaternionFromEuler([0, 0, 0])
        self._robot_id = self.bullet_client.loadURDF(self._robot_urdf, robot_start_pos, robot_start_ori)


    def _get_observation(self, mesh_idx, if_save):
        scene_observation = self._get_scene_observation(mesh_idx, if_save)

        # get the robot state as part of the observation: joint position and velocity
        robot_state_lst = []
        for i in range(self._num_joints):
            joint_pos, joint_vel, _, _ = self.bullet_client.getJointState(self._robot_id, i)
            robot_state_lst.append([joint_pos, joint_vel])
        
        # get the cartesian value of the last link
        last_link_position = self.bullet_client.getLinkState(self._robot_id, 5, 1, 1)[0]
        robot_state_lst = robot_state_lst + [last_link_position]

        observation = {
            'robot_state': robot_state_lst,
            'scene_observation': scene_observation
        }

        return observation

    def _get_scene_observation(self, mesh_idx, if_save):

        # remove the plane to obtain the robot mesh
        visual_data = dict()
        visual_data[self._plane_id] = common.get_body_colors(self.bullet_client, self._plane_id)
        common.set_visible(self.bullet_client, visual_data, visible=False)

        self._scene_tsdf = TSDFVolume(self._volume_bounds, voxel_size=self._voxel_size)

        # surrounding cameras
        scene_center = self._view_bounds.mean(axis=1)
        cam_positions = [[-1.0, -1.0, 0.42],
                         [-1.0, 1.0, 0.42],
                         [1.0, -1.0, 0.42],
                         [1.0, 1.0, 0.42]] # hard
        for cam_id in range(self._num_cam-1):
            cam_position = cam_positions[cam_id]
            cam_lookat = [scene_center[0], scene_center[1], 0.42] #hard
            cam_up_direction = [0, 0, 1]
            if self.debug_cam:
                _ = self.bullet_client.loadURDF(self.debug_cam_sphere_urdf, cam_position, pybullet.getQuaternionFromEuler([0,0,0]), )
            color_image, depth_image, segmentation_mask, cam_pose_matrix = self.get_scene_cam_data(cam_position, cam_lookat, cam_up_direction)
            self._scene_tsdf.integrate(color_image, depth_image, self._scene_cam_intrinsics, cam_pose_matrix, obs_weight=1.)

        # top-down camera
        cam_position = [scene_center[0], scene_center[1], 1.5] # hard
        cam_lookat = [scene_center[0], scene_center[1], 0]
        cam_up_direction = [0, 1, 0]
        if self.debug_cam:
            _ = self.bullet_client.loadURDF(self.debug_cam_sphere_urdf, cam_position, pybullet.getQuaternionFromEuler([0,0,0]), )
        color_image, depth_image, segmentation_mask, cam_pose_matrix = self.get_scene_cam_data(cam_position, cam_lookat, cam_up_direction)
        self._scene_tsdf.integrate(color_image, depth_image, self._scene_cam_intrinsics, cam_pose_matrix, obs_weight=1.)

        # get scene_tsdf(WxHxD) and obstacle_vol(WxHxDx3)
        scene_tsdf, obstacle_vol = self._scene_tsdf.get_volume()

        if if_save:
            verts, faces, norms, colors = self._scene_tsdf.get_mesh()
            fusion.meshwrite(os.path.join(self.save_mesh_folder, f"mesh_{mesh_idx}.ply"), verts, faces, norms, colors)

        scene_observation = {
            'scene_tsdf': scene_tsdf,
            'obstacle_vol': obstacle_vol,
        }

        # bring the plane back
        common.set_visible(self.bullet_client, visual_data, visible=True)

        return scene_observation

    def step(self, action):

        reward = None
        done = None
        info = None

        forces = np.array([5, 5, 5, 5, 5])

        if action['joint'] is None:
            for i in range(len(action['robot'])):
                self.bullet_client.setJointMotorControl2(self._robot_id,
                                                         i,
                                                         controlMode=self._ctrl_mode,
                                                         targetPosition=action['robot'][i],
                                                         force=forces[i])
        else:
            for p_joint in action['joint']:
                self.bullet_client.setJointMotorControl2(self._robot_id,
                                                        p_joint,
                                                        controlMode=self._ctrl_mode,
                                                        targetPosition=action['robot'][p_joint],
                                                        force=forces[p_joint])
            for i in range(len(action['robot'])):
                if i not in action['joint']:
                    self.bullet_client.setJointMotorControl2(self._robot_id,
                                                             i,
                                                             controlMode=self._ctrl_mode,
                                                             targetPosition=0.,
                                                             force=forces[i])

        # wait until the arm becomes stable or timeout (due to unreachable state)
        success_flag = False
        start_time = datetime.datetime.now()
        while True:
            self.bullet_client.stepSimulation()
            if self._if_robot_stable() and self._if_reach_target_joint(action):
                success_flag = True
                break
            passed_time = (datetime.datetime.now() - start_time).total_seconds()
            if passed_time > self._timeout:
                success_flag = False
                break

        # return observation: robot state, scene observation
        if success_flag:
            observation = self._get_observation(mesh_idx=self._n_steps, if_save=True)
            self._saved_robot_state_dict[self._n_steps] = observation['robot_state']
            self._n_steps = self._n_steps + 1
        else:
            observation = None
        return observation, reward, done, info

    def save_robot_state(self):
        with open(os.path.join(self.save_mesh_folder, 'robot_state.json'), 'w') as file:
            json.dump(self._saved_robot_state_dict, file, indent=4)

    def _if_robot_stable(self):
        joint_vel_lst = []
        for i in range(self._num_joints):
            _, joint_vel, _, _ = self.bullet_client.getJointState(self._robot_id, i)
            joint_vel_lst.append(joint_vel)
        if np.abs(np.sum(joint_vel_lst)) <= 0.005:
            return True
        else:
            return False

    def _if_reach_target_joint(self, action):
        for i in range(self._num_joints):
            joint_pos, joint_vel, _, _ = self.bullet_client.getJointState(self._robot_id, i)
            # if all the joints are being controlled.
            if action['joint'] is None:
                joint_diff = joint_pos - action['robot'][i]
                if np.abs(joint_diff) > 1e-3:
                    return False
            # if only specific joints are being controlled.
            else:
                if i in action['joint']:
                    joint_diff = joint_pos - action['robot'][i]
                else:
                    joint_diff = joint_pos - 0.
                if np.abs(joint_diff) > 1e-3:
                    return False
        return True

if __name__ == '__main__':

    env = ArmPybulletSim(gui_enabled=False, num_cam=5)
    N = 10000
    while True:
        env.reset_everything()
        action = {'robot': np.random.uniform(-1, 1, env._num_joints) * np.pi, 'joint': [0, 1, 2, 3]}
        obs, r, done, _ = env.step(action)
        print(env._n_steps)
        if env._n_steps == N:
            break
    env.save_robot_state()