From 5e08d6d334eedf00adb25e2cc545af08bf33984b Mon Sep 17 00:00:00 2001 From: erwincoumans Date: Wed, 14 Feb 2018 15:37:50 -0800 Subject: [PATCH] small test to create / simulate random box stacks --- .../prediction/boxstack_pybullet_sim.py | 156 +++++++++++++ .../prediction/pybullet_sim_gym_env.py | 219 ++++++++++++++++++ .../prediction/test_pybullet_sim_gym_env.py | 60 +++++ 3 files changed, 435 insertions(+) create mode 100644 examples/pybullet/gym/pybullet_envs/prediction/boxstack_pybullet_sim.py create mode 100644 examples/pybullet/gym/pybullet_envs/prediction/pybullet_sim_gym_env.py create mode 100644 examples/pybullet/gym/pybullet_envs/prediction/test_pybullet_sim_gym_env.py diff --git a/examples/pybullet/gym/pybullet_envs/prediction/boxstack_pybullet_sim.py b/examples/pybullet/gym/pybullet_envs/prediction/boxstack_pybullet_sim.py new file mode 100644 index 000000000..42bc3cffa --- /dev/null +++ b/examples/pybullet/gym/pybullet_envs/prediction/boxstack_pybullet_sim.py @@ -0,0 +1,156 @@ +"""This file implements the functionalities of an example simulated system using pybullet. + +""" +import copy +import math +import numpy as np +import os +import pybullet_data +import random + + +class BoxStackPyBulletSim(object): + """The ExamplePyBulletSim class that adds some objects to the scene, steps the sim and return a reward. + + """ + + def __init__(self, + pybullet_client, + urdf_root= pybullet_data.getDataPath(), + time_step=0.01): + """Constructs an example simulation and reset it to the initial states. + + Args: + pybullet_client: The instance of BulletClient to manage different + simulations. + urdf_root: The path to the urdf folder. + time_step: The time step of the simulation. + """ + self._pybullet_client = pybullet_client + self._urdf_root = urdf_root + self.m_actions_taken_since_reset=0 + self.time_step = time_step + self.stateId = -1 + self.Reset(reload_urdf=True) + + + def Reset(self, reload_urdf=False): + """Reset the minitaur to its initial states. + + Args: + reload_urdf: Whether to reload the urdf file. If not, Reset() just place + the minitaur back to its starting position. + """ + self.m_actions_taken_since_reset=0 + xPosRange=0.025 + yPosRange=0.025 + boxHalfExtents = 0.025 + + if reload_urdf: + camInfo = self._pybullet_client.getDebugVisualizerCamera() + cameraDistance=camInfo[10] + print("cameraYaw=",camInfo[8]) + print("cameraPitch=",camInfo[9]) + print("camtarget=",camInfo[11]) + print("projectionMatrix=",camInfo[3]) + self._pybullet_client.resetDebugVisualizerCamera(cameraDistance=0.3, cameraYaw=camInfo[8], cameraPitch=camInfo[9],cameraTargetPosition=camInfo[11]) + + plane = self._pybullet_client.loadURDF("plane.urdf") + texUid = self._pybullet_client.loadTexture("checker_blue.png") + self._pybullet_client.changeVisualShape(plane,-1, textureUniqueId = texUid) + + + self._numObjects=4 #random number? + + + self._cubes=[] + + red=[0.97,0.25,0.25,1] + green=[0.41,0.68,0.31,1] + yellow=[0.92,0.73,0,1] + blue=[0,0.55,0.81,1] + colors=[red,green,yellow,blue] + + for i in range (self._numObjects): + pos=[0,0,boxHalfExtents + i*2*boxHalfExtents] + orn = self._pybullet_client.getQuaternionFromEuler([0,0,0]) + orn=[0,0,0,1] + cube = self._pybullet_client.loadURDF("cube_small.urdf",pos,orn) + self._pybullet_client.changeVisualShape(cube,-1,rgbaColor=colors[i]) + self._cubes.append(cube) + + self._pybullet_client.setGravity(0, 0, -10) + self.stateId = self._pybullet_client.saveState() + else: + if (self.stateId>=0): + self._pybullet_client.restoreState(self.stateId) + index=0 + for i in self._cubes: + posX = random.uniform(-xPosRange,xPosRange) + posY = random.uniform(-yPosRange,yPosRange) + yaw = random.uniform(-math.pi,math.pi) + pos=[posX,posY,boxHalfExtents + index*2*boxHalfExtents] + index+=1 + orn = self._pybullet_client.getQuaternionFromEuler([0,0,yaw]) + self._pybullet_client.resetBasePositionAndOrientation(i,pos,orn) + + def GetActionDimension(self): + """Get the length of the action list. + + Returns: + The length of the action list. + """ + return 4#self.num_motors + + def GetObservationUpperBound(self): + """Get the upper bound of the observation. + + Returns: + The upper bound of an observation. See GetObservation() for the details + of each element of an observation. + """ + upper_bound = np.array([0.0] * self.GetObservationDimension()) + upper_bound[0:] = 1.0 # Quaternion of base orientation. + return upper_bound + + def GetObservationLowerBound(self): + """Get the lower bound of the observation.""" + return -self.GetObservationUpperBound() + + def GetObservationDimension(self): + """Get the length of the observation list. + + Returns: + The length of the observation list. + """ + sz = len(self.GetObservation()) + print("sz=",sz) + return sz + + def GetObservation(self): + """Get the observations of minitaur. + + Returns: + The observation list. observation[0:4] is the base orientation in quaternion form. + """ + observation = [] + for i in self._cubes: + pos,orn=self._pybullet_client.getBasePositionAndOrientation(i) + observation.extend(list(pos)) + observation.extend(list(orn)) + return observation + + def ApplyAction(self, action): + """Set the desired action. + """ + self.m_actions_taken_since_reset+=1 + + + def Termination(self): + return self.m_actions_taken_since_reset>=200 + +def CreateSim(pybullet_client,urdf_root,time_step): + sim = BoxStackPyBulletSim(pybullet_client=pybullet_client, + urdf_root=urdf_root, + time_step=time_step) + return sim diff --git a/examples/pybullet/gym/pybullet_envs/prediction/pybullet_sim_gym_env.py b/examples/pybullet/gym/pybullet_envs/prediction/pybullet_sim_gym_env.py new file mode 100644 index 000000000..84833f88f --- /dev/null +++ b/examples/pybullet/gym/pybullet_envs/prediction/pybullet_sim_gym_env.py @@ -0,0 +1,219 @@ +"""This file implements the gym environment of example PyBullet simulation. + +""" + +import os, inspect +currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +parentdir = os.path.dirname(os.path.dirname(currentdir)) +os.sys.path.insert(0,parentdir) + + +import math +import time +import gym +from gym import spaces +from gym.utils import seeding +import numpy as np +import pybullet +from pybullet_envs.bullet import bullet_client + +from pybullet_envs.prediction import boxstack_pybullet_sim + +import os +import pybullet_data + +from pkg_resources import parse_version + + +class PyBulletSimGymEnv(gym.Env): + """The gym environment to run pybullet simulations. + + + """ + metadata = { + "render.modes": ["human", "rgb_array"], + "video.frames_per_second": 50 + } + + def __init__(self,pybullet_sim_factory = boxstack_pybullet_sim, + render=True, + render_sleep=False, + debug_visualization=True, + hard_reset = False, + render_width=740, + render_height=740, + action_repeat=1, + time_step = 1./240., + num_bullet_solver_iterations=50, + urdf_root=pybullet_data.getDataPath()): + """Initialize the gym environment. + + Args: + urdf_root: The path to the urdf data folder. + """ + self._pybullet_sim_factory = pybullet_sim_factory + self._time_step = time_step + self._urdf_root = urdf_root + self._observation = [] + self._action_repeat=action_repeat + self._num_bullet_solver_iterations = num_bullet_solver_iterations + self._env_step_counter = 0 + self._is_render = render + self._debug_visualization = debug_visualization + self._render_sleep = render_sleep + self._render_width=render_width + self._render_height=render_height + self._cam_dist = .3 + self._cam_yaw = 50 + self._cam_pitch = -35 + self._hard_reset = True + self._last_frame_time = 0.0 + + optionstring='--width={} --height={}'.format(render_width,render_height) + + print("urdf_root=" + self._urdf_root) + + if self._is_render: + self._pybullet_client = bullet_client.BulletClient( + connection_mode=pybullet.GUI, options=optionstring) + else: + self._pybullet_client = bullet_client.BulletClient() + + if (debug_visualization==False): + self._pybullet_client.configureDebugVisualizer(flag=self._pybullet_client.COV_ENABLE_GUI,enable=0) + self._pybullet_client.configureDebugVisualizer(flag=self._pybullet_client.COV_ENABLE_RGB_BUFFER_PREVIEW,enable=0) + self._pybullet_client.configureDebugVisualizer(flag=self._pybullet_client.COV_ENABLE_DEPTH_BUFFER_PREVIEW,enable=0) + self._pybullet_client.configureDebugVisualizer(flag=self._pybullet_client.COV_ENABLE_SEGMENTATION_MARK_PREVIEW,enable=0) + + + self._pybullet_client.setAdditionalSearchPath(urdf_root) + self._seed() + self.reset() + + observation_high = ( + self._example_sim.GetObservationUpperBound()) + observation_low = ( + self._example_sim.GetObservationLowerBound()) + + action_dim = self._example_sim.GetActionDimension() + self._action_bound = 1 + action_high = np.array([self._action_bound] * action_dim) + self.action_space = spaces.Box(-action_high, action_high) + self.observation_space = spaces.Box(observation_low, observation_high) + + self.viewer = None + self._hard_reset = hard_reset # This assignment need to be after reset() + + + def configure(self, args): + self._args = args + + def _reset(self): + if self._hard_reset: + self._pybullet_client.resetSimulation() + + self._pybullet_client.setPhysicsEngineParameter( + numSolverIterations=int(self._num_bullet_solver_iterations)) + self._pybullet_client.setTimeStep(self._time_step) + + self._example_sim = self._pybullet_sim_factory.CreateSim( + pybullet_client=self._pybullet_client, + urdf_root=self._urdf_root, + time_step=self._time_step) + else: + self._example_sim.Reset(reload_urdf=False) + + self._env_step_counter = 0 + + #self._pybullet_client.resetDebugVisualizerCamera( + # self._cam_dist, self._cam_yaw, self._cam_pitch, [0, 0, 0]) + + return self._get_observation() + + def _seed(self, seed=None): + self.np_random, seed = seeding.np_random(seed) + return [seed] + + def _step(self, action): + """Step forward the simulation, given the action. + + Args: + action: the predicted state + + Returns: + observations: The actual state. + reward: The reward for how well the prediction matches the actual state. + done: Whether the episode has ended. + info: A dictionary that stores diagnostic information. + + Raises: + ValueError: The action dimension is not the same as the number of motors. + ValueError: The magnitude of actions is out of bounds. + """ + if self._render_sleep: + # Sleep, otherwise the computation takes less time than real time, + # which will make the visualization like a fast-forward video. + time_spent = time.time() - self._last_frame_time + self._last_frame_time = time.time() + time_to_sleep = self._action_repeat * self._time_step - time_spent + if time_to_sleep > 0: + time.sleep(time_to_sleep) + + #base_pos = self.minitaur.GetBasePosition() + #self._pybullet_client.resetDebugVisualizerCamera( + # self._cam_dist, self._cam_yaw, self._cam_pitch, base_pos) + + + for _ in range(self._action_repeat): + self._example_sim.ApplyAction(action) + self._pybullet_client.stepSimulation() + + self._env_step_counter += 1 + reward = self._reward() + done = self._termination() + return np.array(self._get_observation()), reward, done, {} + + def _render(self, mode="rgb_array", close=False): + if mode != "rgb_array": + return np.array([]) + base_pos = [0,0,0] + view_matrix = self._pybullet_client.computeViewMatrixFromYawPitchRoll( + cameraTargetPosition=base_pos, + distance=self._cam_dist, + yaw=self._cam_yaw, + pitch=self._cam_pitch, + roll=0, + upAxisIndex=2) + proj_matrix = self._pybullet_client.computeProjectionMatrixFOV( + fov=60, aspect=float(self._render_width)/self._render_width, + nearVal=0.01, farVal=100.0) + proj_matrix=[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, -1.0000200271606445, -1.0, 0.0, 0.0, -0.02000020071864128, 0.0] + (_, _, px, _, _) = self._pybullet_client.getCameraImage( + width=self._render_width, height=self._render_height, viewMatrix=view_matrix, + projectionMatrix=proj_matrix, renderer=pybullet.ER_BULLET_HARDWARE_OPENGL) + rgb_array = np.array(px, dtype=np.uint8) + rgb_array = np.reshape(rgb_array, (self._render_height, self._render_width, 4)) + rgb_array = rgb_array[:, :, :3] + return rgb_array + + + + def _termination(self): + terminate=self._example_sim.Termination() + return terminate + + def _reward(self): + reward = 0 + return reward + + + def _get_observation(self): + self._observation = self._example_sim.GetObservation() + return self._observation + + + if parse_version(gym.__version__)>=parse_version('0.9.6'): + render = _render + reset = _reset + seed = _seed + step = _step diff --git a/examples/pybullet/gym/pybullet_envs/prediction/test_pybullet_sim_gym_env.py b/examples/pybullet/gym/pybullet_envs/prediction/test_pybullet_sim_gym_env.py new file mode 100644 index 000000000..23f734e4e --- /dev/null +++ b/examples/pybullet/gym/pybullet_envs/prediction/test_pybullet_sim_gym_env.py @@ -0,0 +1,60 @@ +r"""An example to run of the minitaur gym environment with sine gaits. +""" + +import os +import inspect +currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +parentdir = os.path.dirname(os.path.dirname(currentdir)) +os.sys.path.insert(0,parentdir) + +import math +import numpy as np +from pybullet_envs.prediction import pybullet_sim_gym_env +import argparse +import time + +from pybullet_envs.prediction import boxstack_pybullet_sim +from gym.wrappers.monitoring import video_recorder + +def ResetPoseExample(steps): + """An example that the minitaur stands still using the reset pose.""" + + + environment = pybullet_sim_gym_env.PyBulletSimGymEnv( + pybullet_sim_factory=boxstack_pybullet_sim, + debug_visualization=False, + render=True, action_repeat=30) + action = [math.pi / 2] * 8 + + vid = video_recorder.VideoRecorder(env=environment,path="vid.mp4") + + for _ in range(steps): + print(_) + startsim = time.time() + _, _, done, _ = environment.step(action) + stopsim = time.time() + startrender = time.time() + #environment.render(mode='rgb_array') + vid.capture_frame() + stoprender = time.time() + print ("env.step " , (stopsim - startsim)) + print ("env.render " , (stoprender - startrender)) + if done: + environment.reset() + + + + +def main(): + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument('--env', help='environment ID (0==reset)',type=int, default=0) + args = parser.parse_args() + print("--env=" + str(args.env)) + + if (args.env == 0): + ResetPoseExample(steps = 1000) + +if __name__ == '__main__': + main() + +