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