yfzhoucs/TinyLanguageRobots
1
1"""2@author: Olivier Sigaud3A merge between two sources:4* Adaptation of the MountainCar Environment from the "FAReinforcement" library5of Jose Antonio Martin H. (version 1.0), adapted by 'Tom Schaul, tom@idsia.ch'6and then modified by Arnaud de Broissia7* the gym MountainCar environment8itself from9http://incompleteideas.net/sutton/MountainCar/MountainCar1.cp10permalink: https://perma.cc/6Z2N-PFWC11"""12# apple: https://unsplash.com/images/food/apple13# orange: https://unsplash.com/s/photos/orange14# wood: https://architextures.org/textures/category/wood15import math16from typing import Optional17 18import numpy as np19 20import gym21from gym import spaces22# from gym.envs.classic_control import utils23from gym.error import DependencyNotInstalled24# from gym.utils.renderer import Renderer25import pygame26import scipy27import yaml28from collections import OrderedDict29import copy30 31 32class TinyUR5Env(gym.Env):33 """34 ### Description35 The Mountain Car MDP is a deterministic MDP that consists of a car placed stochastically36 at the bottom of a sinusoidal valley, with the only possible actions being the accelerations37 that can be applied to the car in either direction. The goal of the MDP is to strategically38 accelerate the car to reach the goal state on top of the right hill. There are two versions39 of the mountain car domain in gym: one with discrete actions and one with continuous.40 This version is the one with continuous actions.41 This MDP first appeared in [Andrew Moore's PhD Thesis (1990)](https://www.cl.cam.ac.uk/techreports/UCAM-CL-TR-209.pdf)42 ```43 @TECHREPORT{Moore90efficientmemory-based,44 author = {Andrew William Moore},45 title = {Efficient Memory-based Learning for Robot Control},46 institution = {University of Cambridge},47 year = {1990}48 }49 ```50 ### Observation Space51 The observation is a `ndarray` with shape `(2,)` where the elements correspond to the following:52 | Num | Observation | Min | Max | Unit |53 |-----|--------------------------------------|------|-----|--------------|54 | 0 | position of the car along the x-axis | -Inf | Inf | position (m) |55 | 1 | velocity of the car | -Inf | Inf | position (m) |56 ### Action Space57 The action is a `ndarray` with shape `(1,)`, representing the directional force applied on the car.58 The action is clipped in the range `[-1,1]` and multiplied by a power of 0.0015.59 ### Transition Dynamics:60 Given an action, the mountain car follows the following transition dynamics:61 *velocity<sub>t+1</sub> = velocity<sub>t+1</sub> + force * self.power - 0.0025 * cos(3 * position<sub>t</sub>)*62 *position<sub>t+1</sub> = position<sub>t</sub> + velocity<sub>t+1</sub>*63 where force is the action clipped to the range `[-1,1]` and power is a constant 0.0015.64 The collisions at either end are inelastic with the velocity set to 0 upon collision with the wall.65 The position is clipped to the range [-1.2, 0.6] and velocity is clipped to the range [-0.07, 0.07].66 ### Reward67 A negative reward of *-0.1 * action<sup>2</sup>* is received at each timestep to penalise for68 taking actions of large magnitude. If the mountain car reaches the goal then a positive reward of +10069 is added to the negative reward for that timestep.70 ### Starting State71 The position of the car is assigned a uniform random value in `[-0.6 , -0.4]`.72 The starting velocity of the car is always assigned to 0.73 ### Episode End74 The episode ends if either of the following happens:75 1. Termination: The position of the car is greater than or equal to 0.45 (the goal position on top of the right hill)76 2. Truncation: The length of the episode is 999.77 ### Arguments78 ```79 gym.make('MountainCarContinuous-v0')80 ```81 ### Version History82 * v0: Initial versions release (1.0.0)83 """84 85 metadata = {86 "render_modes": ["human", "rgb_array", "single_rgb_array"],87 "render_fps": 30,88 }89 90 # def __init__(self, yaml_file='config.yaml', render_mode: Optional[str] = None, goal_velocity=0, initializer=None):91 def __init__(self, config, render_mode: Optional[str] = None, goal_velocity=0):92 93 # with open(yaml_file, "r") as stream:94 # try:95 # config = yaml.safe_load(stream)96 # # print(config, type(config))97 # except yaml.YAMLError as exc:98 # print(exc)99 # # exit()100 101 # if initializer is not None:102 # config = initializer.initialize103 self.config = config104 105 self.min_action = -np.pi * 2 - 0.01106 self.max_action = np.pi * 2 + 0.01107 self.min_position = -1.2108 self.max_position = 0.6109 self.max_speed = 0.07110 self.goal_position = (111 0.45 # was 0.5 in gym, 0.45 in Arnaud de Broissia's version112 )113 self.goal_velocity = goal_velocity114 self.power = 0.0015115 116 self.low_state = np.array(117 [self.min_position, -self.max_speed], dtype=np.float32118 )119 self.high_state = np.array(120 [self.max_position, self.max_speed], dtype=np.float32121 )122 123 self.render_mode = render_mode124 # self.renderer = Renderer(self.render_mode, self._render)125 126 self.scale = config['scale']127 self.screen_width = int(config['desk_width'] * self.scale)128 self.screen_height = int(config['desk_height'] * self.scale)129 self.robot_base_xy = [config['robot']['base_x'] * self.scale, config['robot']['base_y'] * self.scale]130 self.tool_center_point = config['robot']['tool_center_point_distance'] * self.scale131 self.tool_img_mid_point = config['robot']['tool_img_mid_point'] * self.scale132 self.screen = None133 self.clock = None134 self.isopen = True135 136 self.action_space = spaces.Box(137 low=self.min_action, high=self.max_action, shape=(4,), dtype=np.float32138 )139 self.observation_space = spaces.Box(140 low=self.min_action, high=self.max_action, shape=(4,), dtype=np.float32141 )142 143 if 'init_joints' not in config:144 self.robot_joints = np.zeros((4,), dtype=np.float32)145 self.robot_joints[0] = -1.57146 self.robot_joints[1] = 1.57147 self.robot_joints[2] = 0148 self.robot_joints_init = np.zeros((4,), dtype=np.float32)149 self.robot_joints_init[0] = -1.57150 self.robot_joints_init[1] = 1.57151 self.robot_joints_init[2] = 0152 else:153 self.robot_joints = np.zeros((4,), dtype=np.float32)154 for i in range(4):155 self.robot_joints[i] = config['init_joints'][i]156 self.robot_joints_init = copy.deepcopy(self.robot_joints) 157 self.lim_length = config['objects']['lim']['length'] * self.scale158 159 self.Kp = 15160 self.dt = 0.003161 162 self.manip_objs = OrderedDict()163 self.env_objs = {}164 for obj in config['objects']:165 # print(obj)166 # print(config['objects'][obj])167 obj_img = pygame.image.load(config['objects'][obj]['image'])168 169 if 'position' not in config['objects'][obj]:170 self.env_objs[obj] = {}171 self.env_objs[obj]['size_xy'] = [config['objects'][obj]['size']['x'] * self.scale, 172 config['objects'][obj]['size']['y'] * self.scale]173 174 self.env_objs[obj]['image'] = \175 pygame.transform.smoothscale(176 obj_img, 177 self.env_objs[obj]['size_xy'])178 179 if 'position' in config['objects'][obj]:180 self.manip_objs[obj] = {}181 self.manip_objs[obj]['size_xy'] = [config['objects'][obj]['size']['x'] * self.scale, 182 config['objects'][obj]['size']['y'] * self.scale]183 self.manip_objs[obj]['pos_xy'] = [config['objects'][obj]['position']['x'] * self.scale, 184 config['objects'][obj]['position']['y'] * self.scale]185 self.manip_objs[obj]['pos_z'] = config['objects'][obj]['position']['z'] * self.scale186 self.manip_objs[obj]['size_z'] = config['objects'][obj]['size']['z'] * self.scale187 self.manip_objs[obj]['orientation'] = 0.188 self.manip_objs[obj]['lru_score'] = 0189 190 self.manip_objs[obj]['image'] = \191 pygame.transform.smoothscale(192 obj_img, 193 self.manip_objs[obj]['size_xy'])194 # exit()195 196 # print(self.manip_objs)197 self.eef_z = 120 * self.scale198 self.grab = None199 self.grab_position = None200 self.grab_position_z = None201 self.grab_orientation = None202 self.highest_lru_score = 0203 204 205 def _eef_(self):206 207 start_x = 0208 start_y = 0209 end_x = self.robot_base_xy[0]210 end_y = self.robot_base_xy[1]211 angle = 0212 for i in range(self.robot_joints.shape[0] - 1):213 if i < 2:214 start_x = end_x215 start_y = end_y216 angle = angle + self.robot_joints[i]217 end_x = start_x + np.sin(angle) * self.lim_length218 end_y = start_y + np.cos(angle) * self.lim_length219 # mid_x = (start_x + end_x) / 2220 # mid_y = (start_y + end_y) / 2221 222 elif i == 2:223 start_x = end_x224 start_y = end_y225 angle = angle + self.robot_joints[i]226 end_x = start_x + np.sin(angle) * self.tool_center_point227 end_y = start_y + np.cos(angle) * self.tool_center_point228 # mid_x = (start_x + end_x) / 2229 # mid_y = (start_y + end_y) / 2230 return np.array([end_x, end_y])231 232 233 def _eef_orientation_(self):234 # print('joints0', self.robot_joints[0], self.robot_joints[1], self.robot_joints[2])235 # print('joints1', float(self.robot_joints[0]), float(self.robot_joints[1]), float(self.robot_joints[2]))236 return float(self.robot_joints[0]) + float(self.robot_joints[1]) + float(self.robot_joints[2])237 # return float(self.robot_joints[0] + self.robot_joints[1] + self.robot_joints[2])238 239 240 def _l2_(self, eef, position):241 return ((eef[0] - position[0]) ** 2 + (eef[1] - position[1]) ** 2) ** (1/2)242 243 244 def _grab_(self, position, eef):245 grab = (self._l2_(eef, position) < 50 * self.scale)246 # print(self._l2_(eef, position))247 return grab248 249 250 def _grab_z_(self, env_obj, eef_z):251 lower_bound = env_obj['pos_z']252 upper_bound = env_obj['pos_z'] + env_obj['size_z']253 # print(f'bounds: {lower_bound}, {upper_bound}')254 if eef_z >= lower_bound and eef_z <= upper_bound:255 return True256 else:257 return False258 259 260 def _gripper_closed_(self):261 if self.robot_joints[-1] >= 0:262 return True263 else:264 return False265 266 267 def _get_serialized_objs_(self):268 objs = OrderedDict()269 keys = [270 'size_xy',271 'pos_xy',272 'pos_z',273 'size_z',274 'orientation',275 'lru_score',276 ]277 for obj in self.manip_objs:278 objs[obj] = {}279 for key in keys:280 objs[obj][key] = self.manip_objs[obj][key]281 return objs282 283 def _max_speed_(self, disp, lim):284 if disp > lim:285 return lim286 if disp < -lim:287 return -lim288 return disp289 290 def step(self, action: np.ndarray, eef_z=None):291 # print([self.manip_objs[env_obj]['lru_score'] for env_obj in self.manip_objs if 'pos_xy' in self.manip_objs[env_obj]])292 # print(action, self.robot_joints)293 # Convert a possible numpy bool to a Python bool.294 terminated = False295 reward = 0296 297 # action is target angles of the joints298 assert action.shape[0] == self.robot_joints.shape[0]299 for i in range(action.shape[0]):300 # print('before change', self.robot_joints[i])301 # self.robot_joints[i] = self.robot_joints[i] + self.Kp * self._ang_diff(action[i], self.robot_joints[i]) * self.dt302 # self.robot_joints[i] = self.robot_joints[i] + self._max_speed_(303 # self._ang_diff(action[i], self.robot_joints[i]), self.Kp * self.dt)304 # print('after change', self.robot_joints[i])305 self.robot_joints[i] = self.robot_joints[i] + self.Kp * 2 * self._ang_diff(action[i], self.robot_joints[i]) * self.dt306 307 308 eef = self._eef_()309 if eef_z is not None:310 self.eef_z = eef_z * self.scale311 # print(f'eef_z: {self.eef_z}')312 313 # print('grab', self.grab)314 if self.grab is not None:315 assert self.grab_orientation is not None316 self.manip_objs[self.grab]['pos_xy'] = eef + self.grab_position317 self.manip_objs[self.grab]['pos_z'] = self.eef_z + self.grab_position_z318 # print('grab orientation', self.grab_orientation)319 # print('eef orientation', self._eef_orientation_())320 self.manip_objs[self.grab]['orientation'] = self._eef_orientation_() + self.grab_orientation321 # print('orientation', self.manip_objs[self.grab]['orientation'])322 323 if self.manip_objs[self.grab]['pos_z'] < 0:324 self.manip_objs[self.grab]['pos_z'] = 0325 self.grab_position_z = self.manip_objs[self.grab]['pos_z'] - self.eef_z326 # print(self.grab_position, eef, self.positions[1], 1)327 328 # self.grab = None329 # self.grab_position = None330 if self._gripper_closed_():331 if self.grab is None:332 for obj in self.manip_objs:333 if 'pos_xy' not in self.manip_objs[obj]:334 continue335 if self._grab_(self.manip_objs[obj]['pos_xy'], eef) and self._grab_z_(self.manip_objs[obj], self.eef_z):336 self.grab = obj337 self.grab_position = self.manip_objs[obj]['pos_xy'] - eef338 self.grab_position_z = self.manip_objs[self.grab]['pos_z'] - self.eef_z339 assert self.manip_objs[self.grab]['orientation'] is not None340 self.grab_orientation = self.manip_objs[self.grab]['orientation'] - self._eef_orientation_()341 self.highest_lru_score += 1342 self.manip_objs[obj]['lru_score'] = self.highest_lru_score343 self.manip_objs.move_to_end(obj, last=False)344 break345 else:346 self.grab = None347 self.grab_position = None348 self.grab_position_z = None349 self.grab_orientation = None350 # print(self.grab_position, eef, self.positions[1], 2)351 352 353 # self.renderer.render_step()354 355 state = {356 'joints': copy.deepcopy(self.robot_joints),357 'eef': copy.deepcopy(self._eef_()),358 'eef_z': copy.deepcopy(self.eef_z),359 'eef_orientation': copy.deepcopy(self._eef_orientation_()),360 'positions': copy.deepcopy(self._get_serialized_objs_()),361 'grabbed_object': copy.deepcopy(self.grab),362 'grab_position': copy.deepcopy(self.grab_position),363 'grab_position_z': copy.deepcopy(self.grab_position_z),364 'grab_orientation': copy.deepcopy(self.grab_orientation),365 'scale': copy.deepcopy(self.scale),366 }367 368 return state, reward, terminated, {}369 370 def reset(371 self,372 *,373 seed: Optional[int] = None,374 return_info: bool = False,375 options: Optional[dict] = None376 ):377 super().reset(seed=seed)378 # Note that if you use custom reset bounds, it may lead to out-of-bound379 # state/observations.380 # low, high = utils.maybe_parse_reset_bounds(options, -0.6, -0.4)381 # self.state = np.array([self.np_random.uniform(low=low, high=high), 0])382 self.state = np.array([self.np_random.uniform(low=-0.6, high=-0.4), 0])383 # self.renderer.reset()384 # self.renderer.render_step()385 if not return_info:386 return np.array(self.state, dtype=np.float32)387 else:388 return np.array(self.state, dtype=np.float32), {}389 390 def _height(self, xs):391 return np.sin(3 * xs) * 0.45 + 0.55392 393 394 def _ang_diff(self, theta1, theta2):395 # Returns the difference between two angles in the range -pi to +pi396 # print('angle diff', (theta1 - theta2 + np.pi) % (2 * np.pi) - np.pi)397 # print(f'theta1 {theta1} theta2 {theta2}')398 return (theta1 - theta2 + np.pi) % (2 * np.pi) - np.pi399 400 def render(self, mode="human"):401 # if self.render_mode is not None:402 # return self.renderer.get_renders()403 # else:404 # return self._render(mode)405 return self._render(mode)406 407 408 def _blitRotate(self, surf, image, origin, pivot, angle):409 image_rect = image.get_rect(topleft = (origin[0] - pivot[0], origin[1]-pivot[1]))410 offset_center_to_pivot = pygame.math.Vector2(origin) - image_rect.center411 rotated_offset = offset_center_to_pivot.rotate(-angle)412 rotated_image_center = (origin[0] - rotated_offset.x, origin[1] - rotated_offset.y)413 rotated_image = pygame.transform.rotate(image, angle)414 rotated_image_rect = rotated_image.get_rect(center = rotated_image_center)415 surf.blit(rotated_image, rotated_image_rect)416 417 418 def get_pos_xy(self, manip_obj):419 return self.manip_objs[manip_obj]['pos_xy']420 421 422 def get_pos_orientation(self, manip_obj):423 return self.manip_objs[manip_obj]['orientation']424 425 426 def ik(self, xyo):427 428 def x_constraint(q, xyo):429 """Returns the corresponding hand xy coordinates for430 a given set of joint angle values [shoulder, elbow, wrist],431 and the above defined arm segment lengths, L432 q : np.array433 the list of current joint angles434 xy : np.array435 current xy position (not used)436 returns : np.array437 the difference between current and desired x position438 """439 xy = xyo[:2]440 xy = [self.scale * x for x in xy]441 return self.lim_length * np.sin(q[0]) + self.lim_length * np.sin(q[0] + q[1]) + self.tool_center_point * np.sin(q[0] + q[1] + q[2]) - xy[0]442 443 def y_constraint(q, xyo):444 """Returns the corresponding hand xy coordinates for445 a given set of joint angle values [shoulder, elbow, wrist],446 and the above defined arm segment lengths, L447 q : np.array448 the list of current joint angles449 xy : np.array450 current xy position (not used)451 returns : np.array452 the difference between current and desired y position453 """454 xy = xyo[:2]455 xy = [self.scale * x for x in xy]456 return self.lim_length * np.cos(q[0]) + self.lim_length * np.cos(q[0] + q[1]) + self.tool_center_point * np.cos(q[0] + q[1] + q[2]) - xy[1]457 458 def sin_o_constraint(q, xyo):459 """Returns the corresponding hand xy coordinates for460 a given set of joint angle values [shoulder, elbow, wrist],461 and the above defined arm segment lengths, L462 q : np.array463 the list of current joint angles464 xy : np.array465 current xy position (not used)466 returns : np.array467 the difference between current and desired y position468 """469 470 o = xyo[-1]471 return (np.sin(q[0] + q[1] + q[2]) - np.sin(o)) ** 2472 473 def cos_o_constraint(q, xyo):474 """Returns the corresponding hand xy coordinates for475 a given set of joint angle values [shoulder, elbow, wrist],476 and the above defined arm segment lengths, L477 q : np.array478 the list of current joint angles479 xy : np.array480 current xy position (not used)481 returns : np.array482 the difference between current and desired y position483 """484 485 o = xyo[-1]486 return (np.cos(q[0] + q[1] + q[2]) - np.cos(o)) ** 2 + (np.sin(q[0] + q[1] + q[2]) - np.sin(o)) ** 2487 488 def distance_to_default(q, *args):489 """Objective function to minimize490 Calculates the euclidean distance through joint space to the491 default arm configuration. The weight list allows the penalty of492 each joint being away from the resting position to be scaled493 differently, such that the arm tries to stay closer to resting494 state more for higher weighted joints than those with a lower495 weight.496 q : np.array497 the list of current joint angles498 returns : scalar499 euclidean distance to the default arm position500 """501 # weights found with trial and error,502 # get some wrist bend, but not much503 weight = [1, 1, 0.5]504 try:505 action = np.sqrt(np.sum([(qi - q0i)**2 * wi506 for qi, q0i, wi in zip(q, self.robot_joints_init.tolist()[:-1], weight)]))507 except Exception:508 print(Exception)509 return action510 511 if len(xyo) == 3:512 ik_result = scipy.optimize.fmin_slsqp(513 func=distance_to_default,514 x0=self.robot_joints,515 eqcons=[x_constraint,516 y_constraint,517 # sin_o_constraint,518 cos_o_constraint],519 # uncomment to add in min / max angles for the joints520 # ieqcons=[joint_limits_upper_constraint,521 # joint_limits_lower_constraint],522 args=(xyo,),523 iprint=0) # iprint=0 suppresses output524 elif len(xyo) == 2:525 ik_result = scipy.optimize.fmin_slsqp(526 func=distance_to_default,527 x0=self.robot_joints,528 eqcons=[x_constraint,529 y_constraint,530 # sin_o_constraint,531 ],532 # uncomment to add in min / max angles for the joints533 # ieqcons=[joint_limits_upper_constraint,534 # joint_limits_lower_constraint],535 args=(xyo,),536 iprint=0) # iprint=0 suppresses output537 return ik_result538 539 def _calculate_img_starting_pos(self, img_pos, img_size):540 x = img_pos[0] - img_size[0] / 2541 y = img_pos[1] - img_size[1] / 2542 return [x, y]543 544 def _render(self, mode="human"):545 assert mode in self.metadata["render_modes"]546 547 try:548 import pygame549 from pygame import gfxdraw550 except ImportError:551 raise DependencyNotInstalled(552 "pygame is not installed, run `pip install gym[classic_control]`"553 )554 555 if self.screen is None:556 pygame.init()557 if mode == "human":558 pygame.display.init()559 self.screen = pygame.display.set_mode(560 (self.screen_width, self.screen_height)561 )562 else: # mode in {"rgb_array", "single_rgb_array"}563 self.screen = pygame.Surface((self.screen_width, self.screen_height))564 if self.clock is None:565 self.clock = pygame.time.Clock()566 567 # world_width = 200568 # scale = self.screen_width / world_width569 570 self.surf = pygame.Surface((self.screen_width, self.screen_height))571 self.surf.fill((255, 255, 255))572 573 574 self.surf.blit(self.env_objs['wood']['image'], (0, 0))575 for obj in list(self.manip_objs.keys())[::-1]:576 if 'pos_xy' in self.manip_objs[obj]:577 578 try:579 image_transformed = pygame.transform.rotate(580 self.manip_objs[obj]['image'], np.rad2deg(self.manip_objs[obj]['orientation']))581 except Exception as e:582 print(e)583 print(self.manip_objs[obj]['orientation'], np.rad2deg(self.manip_objs[obj]['orientation']))584 # print(obj, self.manip_objs[obj]['orientation'])585 new_rect = image_transformed.get_rect()586 # self.surf.blit(image_lim_transformed, (mid_x - new_rect[2] / 2, mid_y - new_rect[3] / 2))587 588 589 self.surf.blit(image_transformed, self._calculate_img_starting_pos(590 self.manip_objs[obj]['pos_xy'],591 [new_rect[2], new_rect[3]]592 ))593 # self.surf.blit(self.image_apple, self.positions[0] - self.size[0] / 2)594 # self.surf.blit(self.image_orange, self.positions[1] - self.size[1] / 2)595 # self.surf.blit(self.image_banana, self.positions[2] - self.size[2] / 2)596 597 start_x = 0598 start_y = 0599 end_x = self.robot_base_xy[0]600 end_y = self.robot_base_xy[1]601 angle = 0602 for i in range(self.robot_joints.shape[0] - 1):603 if i < 2:604 start_x = end_x605 start_y = end_y606 angle = angle + self.robot_joints[i]607 end_x = start_x + np.sin(angle) * self.lim_length608 end_y = start_y + np.cos(angle) * self.lim_length609 mid_x = (start_x + end_x) / 2610 mid_y = (start_y + end_y) / 2611 612 image_lim_transformed = pygame.transform.rotate(self.env_objs['lim']['image'], np.rad2deg(angle))613 new_rect = image_lim_transformed.get_rect()614 self.surf.blit(image_lim_transformed, (mid_x - new_rect[2] / 2, mid_y - new_rect[3] / 2))615 elif i == 2:616 start_x = end_x617 start_y = end_y618 angle = angle + self.robot_joints[i]619 620 # print('gripper angle:', self.robot_joints[-1])621 if self._gripper_closed_():622 mid_x = start_x + np.sin(angle) * (self.env_objs['gripper_closed']['size_xy'][1] / 2 + 38 * self.scale)623 mid_y = start_y + np.cos(angle) * (self.env_objs['gripper_closed']['size_xy'][1] / 2 + 38 * self.scale)624 image_gripper_closed_transformed = pygame.transform.rotate(self.env_objs['gripper_closed']['image'], np.rad2deg(angle))625 new_rect = image_gripper_closed_transformed.get_rect()626 self.surf.blit(image_gripper_closed_transformed, (mid_x - new_rect[2] / 2, mid_y - new_rect[3] / 2))627 else:628 mid_x = start_x + np.sin(angle) * (self.env_objs['gripper_open']['size_xy'][1] / 2 + 38 * self.scale)629 mid_y = start_y + np.cos(angle) * (self.env_objs['gripper_open']['size_xy'][1] / 2 + 38 * self.scale)630 image_gripper_open_transformed = pygame.transform.rotate(self.env_objs['gripper_open']['image'], np.rad2deg(angle))631 new_rect = image_gripper_open_transformed.get_rect()632 self.surf.blit(image_gripper_open_transformed, (mid_x - new_rect[2] / 2, mid_y - new_rect[3] / 2))633 634 # self.surf = pygame.transform.flip(self.surf, False, True)635 self.screen.blit(self.surf, (0, 0))636 if mode == "human":637 pygame.event.pump()638 self.clock.tick(self.metadata["render_fps"])639 pygame.display.flip()640 641 elif mode in {"rgb_array", "single_rgb_array"}:642 return np.transpose(643 np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)644 )645 646 def close(self):647 if self.screen is not None:648 import pygame649 650 pygame.display.quit()651 pygame.quit()652 self.isopen = False653 654 def get_base_xy(self):655 return self.robot_base_xy656 657 def get_joint_angles(self):658 return self.robot_joints659 660 661if __name__ == '__main__':662 # env = TinyUR5Env(render_mode='rgb_array')663 env = TinyUR5Env(render_mode='human')664 665 observation, info = env.reset(seed=42, return_info=True)666 667 for i in range(1000):668 print(i)669 input()670 action = env.action_space.sample()671 observation, reward, done, info = env.step(action)672 img = env.render()673 # print(img[0].shape)674 675 if done:676 observation, info = env.reset(return_info=True)677 env.close()678