small test to create / simulate random box stacks

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

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

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()
+
+