diff --git a/pyrobolearn/__init__.py b/pyrobolearn/__init__.py
new file mode 100644
index 0000000..c9a0d47
--- /dev/null
+++ b/pyrobolearn/__init__.py
@@ -0,0 +1,97 @@
+
+import sys
+
+# import simulators
+# from simulators import *
+import simulators
+
+# import robots
+# from robots import *
+import robots
+
+# import worlds
+# from worlds import *
+import worlds
+
+# import states
+# from states import *
+import states
+
+# import actions
+# from actions import *
+import actions
+
+# import rewards
+# from rewards import *
+import rewards
+
+# import environments
+# from envs import *
+import envs
+
+# import models
+import models
+
+# import approximators
+import approximators
+
+# import policies
+import policies
+
+# import values
+
+# import actor-critics
+
+# import dynamical models
+
+# import tools (interfaces and bridges)
+# import tools
+
+# import tasks
+import tasks
+
+# import metrics
+
+# import algos
+import algos
+
+# import experiments
+
+
+# Meta-information about the package
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "(c) Brian Delhaisse"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+# https://stackoverflow.com/questions/30483246/how-to-check-if-a-python-module-has-been-imported
+# https://stackoverflow.com/questions/14050281/how-to-check-if-a-python-module-exists-without-importing-it/25045228
+def module_imported(module_name): # TODO: improve this method
+ """Check if the given module has been already imported."""
+ if not isinstance(module_name, str):
+ module_name = str(module_name)
+ if module_name in sys.modules:
+ return True
+ return False
+
+
+# Define what submodules/classes/functions should be loaded when writing 'from pyrobolearn import *'
+# __all__ = [
+# # Submodules
+#
+# # Classes
+#
+# # Functions
+#
+# # Context managers
+#
+# # package information
+# "__version__",
+# # Deprecated
+#
+# ]
diff --git a/pyrobolearn/simulators/__init__.py b/pyrobolearn/simulators/__init__.py
new file mode 100644
index 0000000..9586901
--- /dev/null
+++ b/pyrobolearn/simulators/__init__.py
@@ -0,0 +1,19 @@
+
+# load all simulators
+
+# basic simulator
+from simulator import Simulator
+
+# PyBullet simulator
+import pybullet
+import pybullet_data
+from pybullet_envs.bullet.bullet_client import BulletClient
+
+
+def BulletSim(mode=pybullet.GUI, debug_visualizer=False):
+ """mode: pybullet.GUI, pybullet.DIRECT"""
+ sim = BulletClient(connection_mode=mode)
+ sim.setAdditionalSearchPath(pybullet_data.getDataPath())
+ if not debug_visualizer:
+ sim.configureDebugVisualizer(sim.COV_ENABLE_GUI, 0)
+ return sim
diff --git a/pyrobolearn/simulators/bullet.py b/pyrobolearn/simulators/bullet.py
new file mode 100644
index 0000000..175295e
--- /dev/null
+++ b/pyrobolearn/simulators/bullet.py
@@ -0,0 +1,3411 @@
+#!/usr/bin/env python
+"""Define the Bullet Simulator API.
+
+This is the main interface that communicates with the PyBullet simulator [1]. By defining this interface, it allows to
+decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+PyBullet. For instance, some methods in PyBullet do not accepts numpy arrays but only lists. The interface provided
+here makes the necessary conversions.
+
+The signature of each method defined here are inspired by [1,2] but in accordance with the PEP8 style guide [3].
+Parts of the documentation for the methods have been copied-pasted from [2] for completeness purposes.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] PyBullet: https://pybullet.org
+ [2] PyBullet Quickstart Guide: https://docs.google.com/document/d/10sXEhzFRSnvFcl3XxNGhnD4N2SedqwdAvK3dsihxVUA
+ [3] PEP8: https://www.python.org/dev/peps/pep-0008/
+"""
+
+import time
+import numpy as np
+import quaternion
+from pyrobolearn.utils.converter import NumpyListConverter, QuaternionListConverter
+
+import pybullet
+from pybullet_envs.bullet.bullet_client import BulletClient
+
+from simulator import Simulator
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class Bullet(Simulator):
+ r"""PyBullet simulator.
+
+ This is a wrapper around the PyBullet API [1]. For many methods, it is just the same as calling directly the
+ original methods. However for several ones, it converts the data into the correct data type.
+ For instance, some methods in PyBullet returns a matrix :math:`NxM` in a list format with length :math:`NxM`,
+ instead of a numpy array. Other data types includes vectors, quaternions, and others which are all returned as
+ list. The problem with this approach is that we first have to convert the data in our code in order to operate
+ on it. A converter can be specified which converts into the desired format. If none, it will convert the data
+ into numpy arrays instead of lists.
+
+ Also, this wrapper enforces consistency. For instance, all the given and produced angles are represented in
+ radians, and not in degrees. Some original `pybullet` methods require angles expressed in radians, and others in
+ degrees.
+
+ The class also presents the documentation of each method which relieve us to check the user guide [1].
+ Most of the documentation has been copied-pasted from [1], written by Erwin Coumans and Yunfei Bai.
+ Also, Some extra methods have been defined.
+
+ Finally, note that this API is incompatible with the original `pybullet`, i.e. it is not interchangeable in the
+ code! In addition, using this interface allows us to easily switch with other `Simulator` APIs, and make it more
+ modular if the signature of the original PyBullet library change.
+
+ In the following documentation:
+ * `vec3` specifies a list/tuple/np.array of 3 floats
+ * `quat` specifies a list/tuple/np.quaternion of 4 floats
+
+ Examples:
+ sim = Bullet()
+
+ References:
+ [1] PyBullet Quickstart Guide: https://docs.google.com/document/d/10sXEhzFRSnvFcl3XxNGhnD4N2SedqwdAvK3dsihxVUA
+ Erwin Coumans and Yunfei Bai, 2017/2018
+ """
+
+ def __init__(self, render=True): # , converter=None):
+ super(Bullet, self).__init__()
+
+ # Connect to pybullet
+ if render:
+ self.sim = BulletClient(connection_mode=pybullet.GUI)
+ else:
+ self.sim = BulletClient(connection_mode=pybullet.DIRECT)
+ self.id = self.sim._client
+
+ # Converters
+ # if converter is None:
+ self.conv = NumpyListConverter()
+ self.quat_conv = QuaternionListConverter(convention=1)
+
+ # def __del__(self):
+ # """Clean up connection if not already done.
+ #
+ # Copied-pasted from `pybullet_envs/bullet/bullet_client.py`.
+ # """
+ # try:
+ # pybullet.disconnect(physicsClientId=self._client)
+ # except pybullet.error:
+ # pass
+ #
+ # def __getattr__(self, name):
+ # """Inject the client id into Bullet functions.
+ #
+ # Copied-pasted from `pybullet_envs/bullet/bullet_client.py`.
+ # """
+ # attribute = getattr(pybullet, name)
+ # if inspect.isbuiltin(attribute):
+ # attribute = functools.partial(attribute, physicsClientId=self._client)
+ # return attribute
+
+ ##############
+ # Properties #
+ ##############
+
+ @property
+ def version(self):
+ """Return the version of the simulator in a year-month-day format."""
+ return self.sim.getAPIVersion()
+
+ ###########
+ # Methods #
+ ###########
+
+ ##############
+ # Simulators #
+ ##############
+
+ def reset(self):
+ """Reset the simulator.
+
+ "It will remove all objects from the world and reset the world to initial conditions." [1]
+ """
+ self.sim.resetSimulation()
+
+ def step(self, sleep_time=0.):
+ """Perform a step in the simulator.
+
+ "stepSimulation will perform all the actions in a single forward dynamics simulation step such as collision
+ detection, constraint solving and integration. The default timestep is 1/240 second, it can be changed using
+ the setTimeStep or setPhysicsEngineParameter API." [1]
+ """
+ self.sim.stepSimulation()
+ time.sleep(sleep_time)
+
+ def render(self, flag=True):
+ """Render the GUI.
+
+ Args:
+ flag (bool): If True, it will render the simulator by enabling the GUI.
+ """
+ if flag:
+ self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING, 1)
+ else:
+ self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING, 0)
+
+ def set_time_step(self, time_step):
+ """Set the specified time step in the simulator.
+
+ "Warning: in many cases it is best to leave the timeStep to default, which is 240Hz. Several parameters are
+ tuned with this value in mind. For example the number of solver iterations and the error reduction parameters
+ (erp) for contact, friction and non-contact joints are related to the time step. If you change the time step,
+ you may need to re-tune those values accordingly, especially the erp values.
+ You can set the physics engine timestep that is used when calling 'stepSimulation'. It is best to only call
+ this method at the start of a simulation. Don't change this time step regularly. setTimeStep can also be
+ achieved using the new setPhysicsEngineParameter API." [1]
+
+ Args:
+ time_step (float): Each time you call 'step' the time step will proceed with 'time_step'.
+ """
+ self.sim.setTimeStep(timeStep=time_step)
+
+ def set_real_time(self, flag=True):
+ """Enable/disable real time in the simulator.
+
+ "By default, the physics server will not step the simulation, unless you explicitly send a 'stepSimulation'
+ command. This way you can maintain control determinism of the simulation. It is possible to run the simulation
+ in real-time by letting the physics server automatically step the simulation according to its real-time-clock
+ (RTC) using the setRealTimeSimulation command. If you enable the real-time simulation, you don't need to call
+ 'stepSimulation'.
+
+ Note that setRealTimeSimulation has no effect in DIRECT mode: in DIRECT mode the physics server and client
+ happen in the same thread and you trigger every command. In GUI mode and in Virtual Reality mode, and TCP/UDP
+ mode, the physics server runs in a separate thread from the client (PyBullet), and setRealTimeSimulation
+ allows the physicsserver thread to add additional calls to stepSimulation." [1]
+
+ Args:
+ flag (bool): If True, it will enable the real-time simulation. If False, it will disable it.
+ """
+ self.sim.setRealTimeSimulation(enableRealTimeSimulation=int(flag))
+
+ def pause(self):
+ """Pause the simulator if in real-time."""
+ self.set_real_time(False)
+
+ def unpause(self):
+ """Unpause the simulator if in real-time."""
+ self.set_real_time(True)
+
+ def get_physics_properties(self):
+ """Get the physics engine parameters.
+
+ Returns:
+ dict: dictionary containing the following tags with their corresponding values: ['gravityAccelerationX',
+ 'useRealTimeSimulation', 'gravityAccelerationZ', 'numSolverIterations', 'gravityAccelerationY',
+ 'numSubSteps', 'fixedTimeStep']
+ """
+ return self.sim.getPhysicsEngineParameters()
+
+ def set_physics_properties(self, time_step=None, num_solver_iterations=None, use_split_impulse=None,
+ split_impulse_penetration_threshold=None, num_sub_steps=None,
+ collision_filter_mode=None, contact_breaking_threshold=None, max_num_cmd_per_1ms=None,
+ enable_file_caching=None, restitution_velocity_threshold=None, erp=None,
+ contact_erp=None, friction_erp=None, enable_cone_friction=None,
+ deterministic_overlapping_pairs=None, solver_residual_threshold=None):
+ """Set the physics engine parameters.
+
+ Args:
+ time_step (float): See the warning in the `set_time_step` section. Physics engine timestep in
+ fraction of seconds, each time you call `step` simulated time will progress this amount.
+ Same as `set_time_step`. Default to 1./240.
+ num_solver_iterations (int): Choose the maximum number of constraint solver iterations. If the
+ `solver_residual_threshold` is reached, the solver may terminate before the `num_solver_iterations`.
+ Default to 50.
+ use_split_impulse (int): Advanced feature, only when using maximal coordinates: split the positional
+ constraint solving and velocity constraint solving in two stages, to prevent huge penetration recovery
+ forces.
+ split_impulse_penetration_threshold (float): Related to 'useSplitImpulse': if the penetration for a
+ particular contact constraint is less than this specified threshold, no split impulse will happen for
+ that contact.
+ num_sub_steps (int): Subdivide the physics simulation step further by 'numSubSteps'. This will trade
+ performance over accuracy.
+ collision_filter_mode (int): Use 0 for default collision filter: (group A&maskB) AND (groupB&maskA).
+ Use 1 to switch to the OR collision filter: (group A&maskB) OR (groupB&maskA).
+ contact_breaking_threshold (float): Contact points with distance exceeding this threshold are not
+ processed by the LCP solver. In addition, AABBs are extended by this number. Defaults to 0.02 in
+ Bullet 2.x.
+ max_num_cmd_per_1ms (int): Experimental: add 1ms sleep if the number of commands executed exceed this
+ threshold
+ enable_file_caching (bool): Set to 0 to disable file caching, such as .obj wavefront file loading
+ restitution_velocity_threshold (float): If relative velocity is below this threshold, restitution will be
+ zero.
+ erp (float): constraint error reduction parameter (non-contact, non-friction)
+ contact_erp (float): contact error reduction parameter
+ friction_erp (float): friction error reduction parameter (when positional friction anchors are enabled)
+ enable_cone_friction (bool): Set to False to disable implicit cone friction and use pyramid approximation
+ (cone is default)
+ deterministic_overlapping_pairs (bool): Set to True to enable and False to disable sorting of overlapping
+ pairs (backward compatibility setting).
+ solver_residual_threshold (float): velocity threshold, if the maximum velocity-level error for each
+ constraint is below this threshold the solver will terminate (unless the solver hits the
+ numSolverIterations). Default value is 1e-7.
+ """
+ kwargs = {}
+ if time_step is not None:
+ kwargs['fixedTimeStep'] = time_step
+ if num_solver_iterations is not None:
+ kwargs['numSolverIterations'] = num_solver_iterations
+ if use_split_impulse is not None:
+ kwargs['useSplitImpulse'] = use_split_impulse
+ if split_impulse_penetration_threshold is not None:
+ kwargs['splitImpulsePenetrationThreshold'] = split_impulse_penetration_threshold
+ if num_sub_steps is not None:
+ kwargs['numSubSteps'] = num_sub_steps
+ if collision_filter_mode is not None:
+ kwargs['collisionFilterMode'] = collision_filter_mode
+ if contact_breaking_threshold is not None:
+ kwargs['contactBreakingThreshold'] = contact_breaking_threshold
+ if max_num_cmd_per_1ms is not None:
+ kwargs['maxNumCmdPer1ms'] = max_num_cmd_per_1ms
+ if enable_file_caching is not None:
+ kwargs['enableFileCaching'] = enable_file_caching
+ if restitution_velocity_threshold is not None:
+ kwargs['restitutionVelocityThreshold'] = restitution_velocity_threshold
+ if erp is not None:
+ kwargs['erp'] = erp
+ if contact_erp is not None:
+ kwargs['contactERP'] = contact_erp
+ if friction_erp is not None:
+ kwargs['frictionERP'] = friction_erp
+ if enable_cone_friction is not None:
+ kwargs['enableConeFriction'] = int(enable_cone_friction)
+ if deterministic_overlapping_pairs is not None:
+ kwargs['deterministicOverlappingPairs'] = int(deterministic_overlapping_pairs)
+ if solver_residual_threshold is not None:
+ kwargs['solverResidualThreshold'] = solver_residual_threshold
+
+ self.sim.setPhysicsEngineParameter(**kwargs)
+
+ def start_logging(self, logging_type, filename, object_unique_ids, max_log_dof, body_unique_id_A, body_unique_id_B,
+ link_index_A, link_index_B, device_type_filter, log_flags):
+ """
+ Start the logging.
+
+ Args:
+ logging_type (int): There are various types of logging implemented.
+ - STATE_LOGGING_MINITAUR (=0): This will require to load the `quadruped/quadruped.urdf` and object
+ unique id from the quadruped. It logs the timestamp, IMU roll/pitch/yaw, 8 leg motor positions
+ (q0-q7), 8 leg motor torques (u0-u7), the forward speed of the torso and mode (unused in
+ simulation).
+ - STATE_LOGGING_GENERIC_ROBOT (=1): This will log a log of the data of either all objects or selected
+ ones (if `object_unique_ids` is provided).
+ - STATE_LOGGING_VIDEO_MP4 (=3): this will open an MP4 file and start streaming the OpenGL 3D
+ visualizer pixels to the file using an ffmpeg pipe. It will require ffmpeg installed. You can
+ also use avconv (default on Ubuntu), just create a symbolic link so that ffmpeg points to avconv.
+ - STATE_LOGGING_CONTACT_POINTS (=5)
+ - STATE_LOGGING_VR_CONTROLLERS (=2)
+ - STATE_LOGGING_PROFILE_TIMINGS (=6): This will dump a timings file in JSON format that can be opened
+ using Google Chrome about://tracing LOAD.
+ filename (str): file name (absolute or relative path) to store the log file data
+ object_unique_ids (list of int): If left empty, the logger may log every object, otherwise the logger just
+ logs the objects in the object_unique_ids list.
+ max_log_dof (int): Maximum number of joint degrees of freedom to log (excluding the base dofs).
+ This applies to STATE_LOGGING_GENERIC_ROBOT_DATA. Default value is 12. If a robot exceeds the number
+ of dofs, it won't get logged at all.
+ body_unique_id_A (int): Applies to STATE_LOGGING_CONTACT_POINTS (=5). If provided,only log contact points
+ involving body_unique_id_A.
+ body_unique_id_B (int): Applies to STATE_LOGGING_CONTACT_POINTS (=5). If provided,only log contact points
+ involving body_unique_id_B.
+ link_index_A (int): Applies to STATE_LOGGING_CONTACT_POINTS (=5). If provided, only log contact points
+ involving link_index_A for body_unique_id_A.
+ link_index_B (int): Applies to STATE_LOGGING_CONTACT_POINTS (=5). If provided, only log contact points
+ involving link_index_B for body_unique_id_B.
+ device_type_filter (int): deviceTypeFilter allows you to select what VR devices to log:
+ VR_DEVICE_CONTROLLER (=1), VR_DEVICE_HMD (=2), VR_DEVICE_GENERIC_TRACKER (=4) or any combination of
+ them. Applies to STATE_LOGGING_VR_CONTROLLERS (=2). Default values is VR_DEVICE_CONTROLLER (=1).
+ log_flags (int): (upcoming PyBullet 1.3.1). STATE_LOG_JOINT_TORQUES (=3), to log joint torques due to
+ joint motors.
+
+ Returns:
+ int: non-negative logging unique id.
+ """
+ kwargs = {}
+ if object_unique_ids is not None:
+ kwargs['objectUniqueIds'] = object_unique_ids
+ if max_log_dof is not None:
+ kwargs['maxLogDof'] = max_log_dof
+ if body_unique_id_A is not None:
+ kwargs['bodyUniqueIdA'] = body_unique_id_A
+ if body_unique_id_B is not None:
+ kwargs['bodyUniqueIdB'] = body_unique_id_B
+ if link_index_A is not None:
+ kwargs['linkIndexA'] = link_index_A
+ if link_index_B is not None:
+ kwargs['linkIndexB'] = link_index_B
+ if device_type_filter is not None:
+ kwargs['deviceTypeFilter'] = device_type_filter
+ if log_flags is not None:
+ kwargs['logFlags'] = log_flags
+
+ self.sim.startStateLogging(logging_type, filename, **kwargs)
+
+ def stop_logging(self, logger_id):
+ """Stop the logging.
+
+ Args:
+ logger_id (int): unique logger id.
+ """
+ self.sim.stopStateLogging(logger_id)
+
+ def set_gravity(self, gravity=(0, 0, -9.81)):
+ """Set the gravity in the simulator with the given acceleration.
+
+ By default, there is no gravitational force enabled in the simulator.
+
+ Args:
+ gravity (list, tuple of 3 floats): acceleration in the x, y, z directions.
+ """
+ self.sim.setGravity(gravity[0], gravity[1], gravity[2])
+
+ def save(self, filename=None):
+ """
+ Save the state of the simulator.
+
+ Args:
+ filename (None, str): path to file to store the state of the simulator. If None, it will save it in
+ memory instead of the disk.
+
+ Returns:
+ int: unique state id. This id can be used to load the state.
+ """
+ if filename is None:
+ return self.sim.saveState()
+ return self.sim.saveBullet(filename)
+
+ def load(self, state):
+ """
+ Load/Restore the simulator to a previous state.
+
+ Args:
+ state (int, str): unique state id, or path to the file containing the state.
+ """
+ if isinstance(state, int):
+ self.sim.restoreState(stateId=state)
+ elif isinstance(state, str):
+ self.sim.restoreState(fileName=state)
+
+ def load_plugin(self, plugin_path, name):
+ """Load a certain plugin in the simulator.
+
+ Few examples can be found at: https://github.com/bulletphysics/bullet3/tree/master/examples/SharedMemory/plugins
+
+ Args:
+ plugin_path (str): path, location on disk where to find the plugin
+ name (str): postfix name of the plugin that is appended to each API
+
+ Returns:
+ int: unique plugin id. If this id is negative, the plugin is not loaded. Once a plugin is loaded, you can
+ send commands to the plugin using `execute_plugin_commands`
+ """
+ return self.sim.loadPlugin(plugin_path, name)
+
+ def execute_plugin_command(self, plugin_id, *args):
+ """Execute the commands on the specified plugin.
+
+ Args:
+ plugin_id (int): unique plugin id.
+ args (list): list of argument values to be interpreted by the plugin. One can be a string, while the
+ others must be integers or float.
+ """
+ kwargs = {}
+ for arg in args:
+ if isinstance(arg, str):
+ kwargs['textArgument'] = arg
+ elif isinstance(arg, int):
+ kwargs.setdefault('intArgs', []).append(arg)
+ elif isinstance(arg, float):
+ kwargs.setdefault('floatArgs', []).append(arg)
+ self.sim.executePluginCommand(plugin_id, **kwargs)
+
+ def unload_plugin(self, plugin_id):
+ """Unload the specified plugin from the simulator.
+
+ Args:
+ plugin_id (int): unique plugin id.
+ """
+ self.sim.unloadPlugin(plugin_id)
+
+ ######################################
+ # loading URDFs, SDFs, MJCFs, meshes #
+ ######################################
+
+ def load_urdf(self, filename, position=None, orientation=None, use_maximal_coordinates=0,
+ use_fixed_base=0, flags=0, scale=1.0):
+ """Load the given URDF file.
+
+ The loadURDF will send a command to the physics server to load a physics model from a Universal Robot
+ Description File (URDF). The URDF file is used by the ROS project (Robot Operating System) to describe robots
+ and other objects, it was created by the WillowGarage and the Open Source Robotics Foundation (OSRF).
+ Many robots have public URDF files, you can find a description and tutorial here:
+ http://wiki.ros.org/urdf/Tutorials
+
+ Important note:
+ most joints (slider, revolute, continuous) have motors enabled by default that prevent free
+ motion. This is similar to a robot joint with a very high-friction harmonic drive. You should set the joint
+ motor control mode and target settings using `pybullet.setJointMotorControl2`. See the
+ `setJointMotorControl2` API for more information.
+
+ Warning:
+ by default, PyBullet will cache some files to speed up loading. You can disable file caching using
+ `setPhysicsEngineParameter(enableFileCaching=0)`.
+
+ Args:
+ filename (str): a relative or absolute path to the URDF file on the file system of the physics server.
+ position (vec3): create the base of the object at the specified position in world space coordinates [X,Y,Z]
+ orientation (quat): create the base of the object at the specified orientation as world space quaternion
+ [X,Y,Z,W]
+ use_maximal_coordinates (int): Experimental. By default, the joints in the URDF file are created using the
+ reduced coordinate method: the joints are simulated using the Featherstone Articulated Body algorithm
+ (btMultiBody in Bullet 2.x). The useMaximalCoordinates option will create a 6 degree of freedom rigid
+ body for each link, and constraints between those rigid bodies are used to model joints.
+ use_fixed_base (bool): force the base of the loaded object to be static
+ flags (int): URDF_USE_INERTIA_FROM_FILE (val=2): by default, Bullet recomputed the inertia tensor based on
+ mass and volume of the collision shape. If you can provide more accurate inertia tensor, use this flag.
+ URDF_USE_SELF_COLLISION (val=8): by default, Bullet disables self-collision. This flag let's you
+ enable it.
+ You can customize the self-collision behavior using the following flags:
+ * URDF_USE_SELF_COLLISION_EXCLUDE_PARENT (val=16) will discard self-collision between links that
+ are directly connected (parent and child).
+ * URDF_USE_SELF_COLLISION_EXCLUDE_ALL_PARENTS (val=32) will discard self-collisions between a
+ child link and any of its ancestors (parents, parents of parents, up to the base).
+ * URDF_USE_IMPLICIT_CYLINDER (val=128), will use a smooth implicit cylinder. By default, Bullet
+ will tessellate the cylinder into a convex hull.
+ scale (float): scale factor to the URDF model.
+
+ Returns:
+ int (non-negative): unique id associated to the load model.
+ """
+ if position is not None:
+ if isinstance(position, np.ndarray):
+ position = position.ravel().tolist()
+ if orientation is not None:
+ if isinstance(orientation, np.ndarray):
+ orientation = orientation.ravel().tolist()
+ elif isinstance(orientation, quaternion.quaternion):
+ orientation = self.quat_conv.convertFrom(orientation)
+
+ return self.sim.loadURDF(filename, position, orientation, use_maximal_coordinates, int(use_fixed_base), flags,
+ scale)
+
+ def load_sdf(self, filename, scaling=1.):
+ """Load the given SDF file.
+
+ The loadSDF command only extracts some essential parts of the SDF related to the robot models and geometry,
+ and ignores many elements related to cameras, lights and so on.
+
+ Args:
+ filename (str): a relative or absolute path to the SDF file on the file system of the physics server.
+ scaling (float): scale factor for the object
+
+ Returns:
+ list(int): list of object unique id for each object loaded
+ """
+ return self.sim.loadSDF(filename, globalScaling=scaling)
+
+ def load_mjcf(self, filename, scaling=1.):
+ """Load the given MJCF file.
+
+ "The loadMJCF command performs basic import of MuJoCo MJCF xml files, used in OpenAI Gym". [1]
+ It will load all the object described in a MJCF file.
+
+ Args:
+ filename (str): a relative or absolute path to the MJCF file on the file system of the physics server.
+ scaling (float): scale factor for the object
+
+ Returns:
+ list(int): list of object unique id for each object loaded
+ """
+ return self.sim.loadMJCF(filename, globalScaling=scaling)
+
+ def load_mesh(self, filename, position, orientation=(0, 0, 0, 1), mass=1., scale=(1., 1., 1.),
+ color=None, flags=None):
+ """
+ Load a mesh in the world (only available in the simulator).
+
+ Args:
+ filename (str): path to file for the mesh. Currently, only Wavefront .obj. It will create convex hulls
+ for each object (marked as 'o') in the .obj file.
+ position (float[3]): position of the mesh in the Cartesian world space (in meters)
+ orientation (float[4], np.quaternion): orientation of the mesh using quaternion.
+ If np.quaternion then it uses the convention (w,x,y,z). If float[4], it uses the convention (x,y,z,w)
+ mass (float): mass of the mesh (in kg). If mass = 0, it won't move even if there is a collision.
+ scale (float[3]): scale the mesh in the (x,y,z) directions
+ color (int[4]): color of the mesh (by default: white and opaque)
+ flags (int, None): if flag = `sim.GEOM_FORCE_CONCAVE_TRIMESH` (=1), this will create a concave static
+ triangle mesh. This should not be used with dynamic/moving objects, only for static (mass=0) terrain.
+
+ Returns:
+ int: unique id of the mesh in the world
+ """
+ kwargs = {}
+ if flags is not None:
+ kwargs['flags'] = flags
+
+ # create collision shape
+ collision_shape = self.sim.createCollisionShape(pybullet.GEOM_MESH, fileName=filename, meshScale=scale,
+ **kwargs)
+
+ if color is not None:
+ kwargs['rgbaColor'] = color
+
+ # create visual shape
+ visual_shape = self.sim.createVisualShape(pybullet.GEOM_MESH, fileName=filename, meshScale=scale, **kwargs)
+
+ # create body
+ mesh = self.sim.createMultiBody(baseMass=mass,
+ baseCollisionShapeIndex=collision_shape,
+ baseVisualShapeIndex=visual_shape,
+ basePosition=position,
+ baseOrientation=orientation)
+
+ return mesh
+
+ ##########
+ # Bodies #
+ ##########
+
+ # TODO: add the other arguments
+ def create_body(self, visual_shape_id=-1, collision_shape_id=-1, mass=0, position=(0., 0., 0.),
+ orientation=(0., 0., 0., 1.)):
+ """Create a body in the simulator.
+
+ Args:
+ visual_shape_id (int): unique id from createVisualShape or -1. You can reuse the visual shape (instancing)
+ collision_shape_id (int): unique id from createCollisionShape or -1. You can re-use the collision shape
+ for multiple multibodies (instancing)
+ mass (int): mass of the base, in kg (if using SI units)
+ position (np.float[3]): Cartesian world position of the base
+ orientation (np.float[4]): Orientation of base as quaternion [x,y,z,w]
+
+ Returns:
+ int: non-negative unique id or -1 for failure.
+ """
+ if isinstance(position, np.ndarray):
+ position = position.ravel().tolist()
+ if isinstance(orientation, np.ndarray):
+ orientation = orientation.ravel().tolist()
+ elif isinstance(orientation, quaternion.quaternion):
+ orientation = self.quat_conv.convertFrom(orientation)
+ return self.sim.createMultiBody(baseMass=mass, baseCollisionShapeIndex=collision_shape_id,
+ baseVisualShapeIndex=visual_shape_id, basePosition=position,
+ baseOrientation=orientation)
+
+ def remove_body(self, body_id):
+ """Remove a particular body in the simulator.
+
+ Args:
+ body_id (int): unique body id.
+ """
+ self.sim.removeBody(body_id)
+
+ def num_bodies(self):
+ """Return the number of bodies present in the simulator.
+
+ Returns:
+ int: number of bodies
+ """
+ return self.sim.getNumBodies()
+
+ def get_body_info(self, body_id):
+ """Get the specified body information.
+
+ Specifically, it returns the base name extracted from the URDF, SDF, MJCF, or other file.
+
+ Args:
+ body_id (int): unique body id.
+
+ Returns:
+ str: base name
+ """
+ return self.sim.getBodyInfo(body_id)
+
+ def get_body_id(self, index):
+ """
+ Get the body id associated to the index which is between 0 and `num_bodies()`.
+
+ Args:
+ index (int): index between [0, `num_bodies()`]
+
+ Returns:
+ int: unique body id.
+ """
+ return self.sim.getBodyUniqueId(index)
+
+ ###############
+ # constraints #
+ ###############
+
+ def create_constraint(self, parent_body_id, parent_link_id, child_body_id, child_link_id, joint_type,
+ joint_axis, parent_frame_position, child_frame_position,
+ parent_frame_orientation=(0., 0., 0., 1.), child_frame_orientation=(0., 0., 0., 1.)):
+ """
+ Create a constaint.
+
+ "URDF, SDF and MJCF specify articulated bodies as a tree-structures without loops. The 'createConstraint'
+ allows you to connect specific links of bodies to close those loops. In addition, you can create arbitrary
+ constraints between objects, and between an object and a specific world frame.
+ It can also be used to control the motion of physics objects, driven by animated frames, such as a VR
+ controller. It is better to use constraints, instead of setting the position or velocity directly for
+ such purpose, since those constraints are solved together with other dynamics constraints." [1]
+
+ Args:
+ parent_body_id (int): parent body unique id
+ parent_link_id (int): parent link index (or -1 for the base)
+ child_body_id (int): child body unique id, or -1 for no body (specify a non-dynamic child frame in world
+ coordinates)
+ child_link_id (int): child link index, or -1 for the base
+ joint_type (int): joint type: JOINT_PRISMATIC (=1), JOINT_FIXED (=4), JOINT_POINT2POINT (=5),
+ JOINT_GEAR (=6)
+ joint_axis (np.float[3]): joint axis, in child link frame
+ parent_frame_position (np.float[3]): position of the joint frame relative to parent CoM frame.
+ child_frame_position (np.float[3]): position of the joint frame relative to a given child CoM frame (or
+ world origin if no child specified)
+ parent_frame_orientation (np.float[4]): the orientation of the joint frame relative to parent CoM
+ coordinate frame
+ child_frame_orientation (np.float[4]): the orientation of the joint frame relative to the child CoM
+ coordinate frame (or world origin frame if no child specified)
+
+ Examples:
+ - `pybullet/examples/quadruped.py`
+ - `pybullet/examples/constraint.py`
+
+ Returns:
+ int: constraint unique id.
+ """
+ return self.sim.createConstraint(parent_body_id, parent_link_id, child_body_id, child_link_id, joint_type,
+ joint_axis, parent_frame_position, child_frame_position,
+ parent_frame_orientation, child_frame_orientation)
+
+ def remove_constraint(self, constraint_id):
+ """
+ Remove the specified constraint.
+
+ Args:
+ constraint_id (int): constraint unique id.
+ """
+ self.sim.removeConstraint(constraint_id)
+
+ def change_constraint(self, constraint_id, child_joint_pivot=None, child_frame_orientation=None, max_force=None,
+ gear_ratio=None, gear_auxiliary_link=None, relative_position_target=None, erp=None):
+ """
+ Change the parameters of an existing constraint.
+
+ Args:
+ constraint_id (int): constraint unique id.
+ child_joint_pivot (np.float[3]): updated position of the joint frame relative to a given child CoM frame
+ (or world origin if no child specified)
+ child_frame_orientation (np.float[4]): updated child frame orientation as quaternion [x,y,z,w]
+ max_force (float): maximum force that constraint can apply
+ gear_ratio (float): the ratio between the rates at which the two gears rotate
+ gear_auxiliary_link (int): In some cases, such as a differential drive, a third (auxilary) link is used as
+ reference pose. See `racecar_differential.py`
+ relative_position_target (float): the relative position target offset between two gears
+ erp (float): constraint error reduction parameter
+ """
+ kwargs = {}
+ if child_joint_pivot is not None:
+ kwargs['jointChildPivot'] = child_joint_pivot
+ if child_frame_orientation is not None:
+ kwargs['jointChildFrameOrientation'] = child_frame_orientation
+ if max_force is not None:
+ kwargs['maxForce'] = max_force
+ if gear_ratio is not None:
+ kwargs['gearRatio'] = gear_ratio
+ if gear_auxiliary_link is not None:
+ kwargs['gearAuxLink'] = gear_auxiliary_link
+ if relative_position_target is not None:
+ kwargs['relativePositionTarget'] = relative_position_target
+ if erp is not None:
+ kwargs['erp'] = erp
+
+ self.sim.changeConstraint(constraint_id, **kwargs)
+
+ def num_constraints(self):
+ """
+ Get the number of constraints created.
+
+ Returns:
+ int: number of constraints created.
+ """
+ return self.sim.getNumConstraints()
+
+ def get_constraint_id(self, index):
+ """
+ Get the constraint unique id associated with the index which is between 0 and `num_constraints()`.
+
+ Args:
+ index (int): index between [0, `num_constraints()`]
+
+ Returns:
+ int: constraint unique id.
+ """
+ return self.sim.getConstraintUniqueId(index)
+
+ def get_constraint_info(self, constraint_id):
+ """
+ Get information about the given constaint id.
+
+ Args:
+ constraint_id (int): constraint unique id.
+
+ Returns:
+ int: parent_body_id
+ int: parent_joint_id (if -1, it is the base)
+ int: child_body_id (if -1, no body; specify a non-dynamic child frame in world coordinates)
+ int: child_link_id (if -1, it is the base)
+ int: constraint/joint type
+ np.float[3]: joint axis
+ np.float[3]: joint pivot (position) in parent CoM frame
+ np.float[3]: joint pivot (position) in specified child CoM frame (or world frame if no specified child)
+ np.float[4]: joint frame orientation relative to parent CoM coordinate frame
+ np.float[4]: joint frame orientation relative to child CoM frame (or world frame if no specified child)
+ float: maximum force that constraint can apply
+ """
+ return self.sim.getConstraintInfo(constraint_id)
+
+ def get_constraint_state(self, constraint_id):
+ """
+ Get the state of the given constraint.
+
+ Args:
+ constraint_id (int): constraint unique id.
+
+ Returns:
+ np.float[D]: applied constraint forces. Its dimension is the degrees of freedom that are affected by
+ the constraint (a fixed constraint affects 6 DoF for example)
+ """
+ return self.sim.getConstraintState(constraint_id)
+
+ ###########
+ # objects #
+ ###########
+
+ def get_mass(self, body_id):
+ """
+ Return the total mass of the robot (=sum of all mass links).
+
+ Args:
+ body_id (int): unique object id, as returned from `load_urdf`.
+
+ Returns:
+ float: total mass of the robot [kg]
+ """
+ return np.sum(self.get_link_masses(body_id, [-1] + list(range(self.num_links(body_id)))))
+
+ def get_base_mass(self, body_id):
+ """Return the base mass of the robot."""
+ return self.get_link_masses(body_id, -1)
+
+ def get_base_name(self, body_id):
+ """
+ Return the base name.
+
+ Args:
+ body_id (int): unique object id.
+
+ Returns:
+ str: base name
+ """
+ return self.sim.getBodyInfo(body_id)[0]
+
+ def get_center_of_mass(self, body_id, link_ids=None):
+ """
+ Return the center of mass position.
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (list of int): link ids associated with the given body id. If None, it will take all the links
+ of the specified body.
+
+ Returns:
+ np.float[3]: center of mass position in the Cartesian world coordinates
+ """
+ if link_ids is None:
+ link_ids = list(range(self.num_links(body_id)))
+
+ pos = self.get_link_world_positions(body_id, link_ids)
+ mass = self.get_link_masses(body_id, link_ids)
+
+ com = np.sum(pos.T * mass, axis=1) / np.sum(mass)
+
+ return com
+
+ def get_base_pose(self, body_id):
+ """
+ Get the current position and orientation of the base (or root link) of the body in Cartesian world coordinates.
+
+ Args:
+ body_id (int): object unique id, as returned from `load_urdf`.
+
+ Returns:
+ np.float[3]: base position
+ np.float[4]: base orientation (quaternion [x,y,z,w])
+ """
+ pos, orientation = self.sim.getBasePositionAndOrientation(body_id)
+ return np.array(pos), np.array(orientation)
+
+ def get_base_position(self, body_id):
+ """
+ Return the base position of the specified body.
+
+ Args:
+ body_id (int): object unique id, as returned from `load_urdf`.
+
+ Returns:
+ np.float[3]: base position.
+ """
+ return self.get_base_pose(body_id)[0]
+
+ def get_base_orientation(self, body_id):
+ """
+ Get the base orientation of the specified body.
+
+ Args:
+ body_id (int): object unique id, as returned from `load_urdf`.
+
+ Returns:
+ np.float[4]: base orientation in the form of a quaternion (x,y,z,w)
+ """
+ return self.get_base_pose(body_id)[1]
+
+ def reset_base_pose(self, body_id, position, orientation):
+ """
+ Reset the base position and orientation of the specified object id.
+
+ "It is best only to do this at the start, and not during a running simulation, since the command will override
+ the effect of all physics simulation. The linear and angular velocity is set to zero. You can use
+ `reset_base_velocity` to reset to a non-zero linear and/or angular velocity." [1]
+
+ Args:
+ body_id (int): unique object id.
+ position (np.float[3]): new base position.
+ orientation (np.float[4]): new base orientation (expressed as a quaternion [x,y,z,w])
+ """
+ self.sim.resetBasePositionAndOrientation(body_id, position, orientation)
+
+ def reset_base_position(self, body_id, position):
+ """
+ Reset the base position of the specified body/object id while preserving its orientation.
+
+ Args:
+ body_id (int): unique object id.
+ position (np.float[3]): new base position.
+ """
+ orientation = self.get_base_orientation(body_id)
+ self.reset_base_pose(body_id, position, orientation)
+
+ def reset_base_orientation(self, body_id, orientation):
+ """
+ Reset the base orientation of the specified body/object id while preserving its position.
+
+ Args:
+ body_id (int): unique object id.
+ orientation (np.float[4]): new base orientation (expressed as a quaternion [x,y,z,w])
+ """
+ position = self.get_base_position(body_id)
+ self.reset_base_pose(body_id, position, orientation)
+
+ def get_base_velocity(self, body_id):
+ """
+ Return the base linear and angular velocities.
+
+ Args:
+ body_id (int): object unique id, as returned from `load_urdf`.
+
+ Returns:
+ np.float[3]: linear velocity of the base in Cartesian world space coordinates
+ np.float[3]: angular velocity of the base in Cartesian world space coordinates
+ """
+ lin_vel, ang_vel = self.sim.getBaseVelocity(body_id)
+ return np.array(lin_vel), np.array(ang_vel)
+
+ def get_base_linear_velocity(self, body_id):
+ """
+ Return the linear velocity of the base.
+
+ Args:
+ body_id (int): object unique id, as returned from `load_urdf`.
+
+ Returns:
+ np.float[3]: linear velocity of the base in Cartesian world space coordinates
+ """
+ return self.get_base_velocity(body_id)[0]
+
+ def get_base_angular_velocity(self, body_id):
+ """
+ Return the angular velocity of the base.
+
+ Args:
+ body_id (int): object unique id, as returned from `load_urdf`.
+
+ Returns:
+ np.float[3]: angular velocity of the base in Cartesian world space coordinates
+ """
+ return self.get_base_velocity(body_id)[1]
+
+ def reset_base_velocity(self, body_id, linear_velocity=None, angular_velocity=None):
+ """
+ Reset the base velocity.
+
+ Args:
+ body_id (int): unique object id.
+ linear_velocity (np.float[3]): new linear velocity of the base.
+ angular_velocity (np.float[3]): new angular velocity of the base.
+ """
+ if linear_velocity is not None and angular_velocity is not None:
+ self.sim.resetBaseVelocity(body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity)
+ elif linear_velocity is not None:
+ self.sim.resetBaseVelocity(body_id, linearVelocity=linear_velocity)
+ elif angular_velocity is not None:
+ self.sim.resetBaseVelocity(body_id, angularVelocity=angular_velocity)
+
+ def reset_base_linear_velocity(self, body_id, linear_velocity):
+ """
+ Reset the base linear velocity.
+
+ Args:
+ body_id (int): unique object id.
+ linear_velocity (np.float[3]): new linear velocity of the base
+ """
+ self.sim.resetBaseVelocity(body_id, linearVelocity=linear_velocity)
+
+ def reset_base_angular_velocity(self, body_id, angular_velocity):
+ """
+ Reset the base angular velocity.
+
+ Args:
+ body_id (int): unique object id.
+ angular_velocity (np.float[3]): new angular velocity of the base
+ """
+ self.sim.resetBaseVelocity(body_id, angularVelocity=angular_velocity)
+
+ def apply_external_force(self, body_id, link_id=-1, force=(0., 0., 0.), position=(0., 0., 0.),
+ flags=pybullet.LINK_FRAME):
+ """
+ Apply the specified external force on the specified position on the body / link.
+
+ "This method will only work when explicitly stepping the simulation using stepSimulation, in other words:
+ setRealTimeSimulation(0). After each simulation step, the external forces are cleared to zero. If you are
+ using 'setRealTimeSimulation(1), applyExternalForce/Torque will have undefined behavior (either 0, 1 or
+ multiple force/torque applications)" [1]
+
+ Args:
+ body_id (int): unique body id.
+ link_id (int): unique link id. If -1, it will be the base.
+ force (np.float[3]): external force to be applied.
+ position (np.float[3]): position on the link where the force is applied. See `flags` for coordinate
+ systems.
+ flags (int): Specify the coordinate system of force/position: either `pybullet.WORLD_FRAME` (=2) for
+ Cartesian world coordinates or `pybullet.LINK_FRAME` (=1) for local link coordinates.
+ """
+ self.sim.applyExternalForce(body_id, link_id, force, position, flags)
+
+ def apply_external_torque(self, body_id, link_id=-1, torque=(0., 0., 0.)):
+ """
+ Apply an external torque on a body, or a link of the body. Note that after each simulation step, the external
+ torques are cleared to 0.
+
+ Warnings: This does not work when using `sim.setRealTimeSimulation(1)`.
+
+ Args:
+ body_id (int): unique body id.
+ link_id (int): link id to apply the torque, if -1 it will apply the torque on the base
+ torque (float[3]): Cartesian torques to be applied on the body
+ """
+ self.sim.applyExternalTorque(body_id, link_id, torque)
+
+ ###################
+ # transformations #
+ ###################
+
+ #############################
+ # robots (joints and links) #
+ #############################
+
+ def num_joints(self, body_id):
+ """
+ Return the total number of joints of the specified body. This is the same as calling `num_links`.
+
+ Args:
+ body_id (int): unique body id.
+
+ Returns:
+ int: number of joints with the associated body id.
+ """
+ return self.sim.getNumJoints(body_id)
+
+ def num_links(self, body_id):
+ """
+ Return the total number of links of the specified body. This is the same as calling `num_joints`.
+
+ Args:
+ body_id (int): unique body id.
+
+ Returns:
+ int: number of links with the associated body id.
+ """
+ return self.num_joints(body_id)
+
+ def get_joint_info(self, body_id, joint_id):
+ """
+ Return information about the given joint about the specified body.
+
+ Note that this method returns a lot of information, so specific methods have been implemented that return
+ only the desired information. Also, note that we do not convert the data here.
+
+ Args:
+ body_id (int): unique body id.
+ joint_id (int): joint id is included in [0..`num_joints(body_id)`].
+
+ Returns:
+ [0] int: the same joint id as the input parameter
+ [1] str: name of the joint (as specified in the URDF/SDF/etc file)
+ [2] int: type of the joint which implie the number of position and velocity variables.
+ The types include JOINT_REVOLUTE (=0), JOINT_PRISMATIC (=1), JOINT_SPHERICAL (=2),
+ JOINT_PLANAR (=3), and JOINT_FIXED (=4).
+ [3] int: q index - the first position index in the positional state variables for this body
+ [4] int: dq index - the first velocity index in the velocity state variables for this body
+ [5] int: flags (reserved)
+ [6] float: the joint damping value (as specified in the URDF file)
+ [7] float: the joint friction value (as specified in the URDF file)
+ [8] float: the positional lower limit for slider and revolute joints
+ [9] float: the positional upper limit for slider and revolute joints
+ [10] float: maximum force specified in URDF. Note that this value is not automatically used.
+ You can use maxForce in 'setJointMotorControl2'.
+ [11] float: maximum velocity specified in URDF. Note that this value is not used in actual
+ motor control commands at the moment.
+ [12] str: name of the link (as specified in the URDF/SDF/etc file)
+ [13] np.float[3]: joint axis in local frame (ignored for JOINT_FIXED)
+ [14] np.float[3]: joint position in parent frame
+ [15] np.float[4]: joint orientation in parent frame
+ [16] int: parent link index, -1 for base
+ """
+ return self.sim.getJointInfo(body_id, joint_id)
+
+ def get_joint_state(self, body_id, joint_id):
+ """
+ Get the joint state.
+
+ Args:
+ body_id (int): body unique id as returned by `load_urdf`, etc.
+ joint_id (int): joint index in range [0..num_joints(body_id)]
+
+ Returns:
+ float: The position value of this joint.
+ float: The velocity value of this joint.
+ np.float[6]: These are the joint reaction forces, if a torque sensor is enabled for this joint it is
+ [Fx, Fy, Fz, Mx, My, Mz]. Without torque sensor, it is [0, 0, 0, 0, 0, 0].
+ float: This is the motor torque applied during the last stepSimulation. Note that this only applies in
+ VELOCITY_CONTROL and POSITION_CONTROL. If you use TORQUE_CONTROL then the applied joint motor torque
+ is exactly what you provide, so there is no need to report it separately.
+ """
+ pos, vel, forces, torque = self.sim.getJointState(body_id, joint_id)
+ return pos, vel, np.array(forces), torque
+
+ def get_joint_states(self, body_id, joint_ids):
+ """
+ Get the joint state of the specified joints.
+
+ Args:
+ body_id (int): body unique id.
+ joint_ids (list of int): list of joint ids.
+
+ Returns:
+ list:
+ float: The position value of this joint.
+ float: The velocity value of this joint.
+ np.float[6]: These are the joint reaction forces, if a torque sensor is enabled for this joint it is
+ [Fx, Fy, Fz, Mx, My, Mz]. Without torque sensor, it is [0, 0, 0, 0, 0, 0].
+ float: This is the motor torque applied during the last `step`. Note that this only applies in
+ VELOCITY_CONTROL and POSITION_CONTROL. If you use TORQUE_CONTROL then the applied joint motor
+ torque is exactly what you provide, so there is no need to report it separately.
+ """
+ return self.sim.getJointStates(body_id, joint_ids)
+
+ def reset_joint_state(self, body_id, joint_id, target_position, target_velocity=0.):
+ """
+ Reset the state of the joint. It is best only to do this at the start, while not running the simulation:
+ `reset_joint_state` overrides all physics simulation. Note that we only support 1-DOF motorized joints at
+ the moment, sliding joint or revolute joints.
+
+ Args:
+ body_id (int): body unique id as returned by `load_urdf`, etc.
+ joint_id (int): joint index in range [0..num_joints(body_id)]
+ target_position (float): the joint position (angle in radians [rad] or position [m])
+ target_velocity (float): the joint velocity (angular [rad/s] or linear velocity [m/s])
+ """
+ self.sim.resetJointState(body_id, joint_id, target_position, target_velocity)
+
+ def enable_joint_force_torque_sensor(self, body_id, joint_id, enable=True):
+ """
+ You can enable or disable a joint force/torque sensor in each joint. Once enabled, if you perform a
+ `step`, the 'get_joint_state' will report the joint reaction forces in the fixed degrees of freedom: a fixed
+ joint will measure all 6DOF joint forces/torques. A revolute/hinge joint force/torque sensor will measure
+ 5DOF reaction forces along all axis except the hinge axis. The applied force by a joint motor is available
+ in the `applied_joint_motor_torque` of `get_joint_state`.
+
+ Args:
+ body_id (int): body unique id as returned by `load_urdf`, etc.
+ joint_id (int): joint index in range [0..num_joints(body_id)]
+ enable (bool): True to enable, False to disable the force/torque sensor
+ """
+ self.sim.enableJointForceTorqueSensor(body_id, joint_id, enable)
+
+ def set_joint_motor_control(self, body_id, joint_id, control_mode=pybullet.POSITION_CONTROL, position=None,
+ velocity=None, force=None, kp=None, kd=None, max_velocity=None):
+ """
+ Set the joint motor control.
+
+ In position control:
+ .. math:: error = Kp (x_{des} - x) + Kd (\dot{x}_{des} - \dot{x})
+
+ In velocity control:
+ .. math:: error = \dot{x}_{des} - \dot{x}
+
+ Note that the maximum forces and velocities are not automatically used for the different control schemes.
+
+ "We can control a robot by setting a desired control mode for one or more joint motors. During the `step`,
+ the physics engine will simulate the motors to reach the given target value that can be reached within
+ the maximum motor forces and other constraints. Each revolute joint and prismatic joint is motorized
+ by default. There are 3 different motor control modes: position control, velocity control and torque control.
+
+ You can effectively disable the motor by using a force of 0. You need to disable motor in order to use direct
+ torque control: `set_joint_motor_control(body_id, joint_id, control_mode=pybullet.VELOCITY_CONTROL,
+ force=force)`"
+
+ Args:
+ body_id (int): body unique id.
+ joint_id (int): joint/link id.
+ control_mode (int): POSITION_CONTROL (=2) (which is in fact CONTROL_MODE_POSITION_VELOCITY_PD),
+ VELOCITY_CONTROL (=0), TORQUE_CONTROL (=1) and PD_CONTROL (=3).
+ position (float): target joint position (used in POSITION_CONTROL).
+ velocity (float): target joint velocity. In VELOCITY_CONTROL and POSITION_CONTROL, the target velocity is
+ the desired velocity of the joint. Note that the target velocity is not the maximum joint velocity.
+ In PD_CONTROL and POSITION_CONTROL/CONTROL_MODE_POSITION_VELOCITY_PD, the final target velocity is
+ computed using:
+ `kp*(erp*(desiredPosition-currentPosition)/dt)+currentVelocity+kd*(m_desiredVelocity - currentVelocity)`
+ force (float): in POSITION_CONTROL and VELOCITY_CONTROL, this is the maximum motor force used to reach the
+ target value. In TORQUE_CONTROL this is the force/torque to be applied each simulation step.
+ kp (float): position (stiffness) gain (used in POSITION_CONTROL).
+ kd (float): velocity (damping) gain (used in POSITION_CONTROL).
+ max_velocity (float): in POSITION_CONTROL this limits the velocity to a maximum.
+ """
+ kwargs = {}
+ if position is not None:
+ kwargs['targetPosition'] = position
+ if velocity is not None:
+ kwargs['targetVelocity'] = velocity
+ if force is not None:
+ kwargs['force'] = force
+ if kp is not None:
+ kwargs['positionGain'] = kp
+ if kd is not None:
+ kwargs['velocityGain'] = kd
+ if max_velocity is not None:
+ kwargs['maxVelocity'] = max_velocity
+ self.sim.setJointMotorControl2(body_id, joint_id, controlMode=control_mode, **kwargs)
+
+ def set_joint_motor_control_array(self, body_id, joint_ids, control_mode=pybullet.POSITION_CONTROL, positions=None,
+ velocities=None, forces=None, kps=None, kds=None):
+ """
+ Instead of making individual calls for each joint, you can pass arrays for all inputs to reduce calling
+ overhead dramatically.
+
+ Args:
+ body_id (int): body unique id.
+ joint_ids (list of int): list of joint id.
+ control_mode (int): POSITION_CONTROL (=2) (which is in fact CONTROL_MODE_POSITION_VELOCITY_PD),
+ VELOCITY_CONTROL (=0), TORQUE_CONTROL (=1) and PD_CONTROL (=3).
+ positions (list of float): list of target joint positions (used in POSITION_CONTROL) the target value is target position of the joint.
+ velocities (list of float): list of target joint velocities (used in PD_CONTROL, VELOCITY_CONTROL and
+ POSITION_CONTROL).
+ forces (list of float): list of forces. In POSITION_CONTROL and VELOCITY_CONTROL, these are the maximum
+ motor forces used to reach the target values. In TORQUE_CONTROL these are the forces/torques to be
+ applied each simulation step.
+ kps (list of float): list of position (stiffness) gains (used in POSITION_CONTROL).
+ kds (list of float): list of velocity (damping) gains (used in POSITION_CONTROL).
+ """
+ kwargs = {}
+ if positions is not None:
+ kwargs['targetPositions'] = positions
+ if velocities is not None:
+ kwargs['targetVelocities'] = velocities
+ if forces is not None:
+ kwargs['forces'] = forces
+ if kps is not None:
+ kwargs['positionGains'] = kps
+ if kds is not None:
+ kwargs['velocityGains'] = kds
+ self.sim.setJointMotorControlArray(body_id, joint_ids, controlMode=control_mode, **kwargs)
+
+ def get_link_state(self, body_id, link_id, compute_velocity=False, compute_forward_kinematics=False):
+ """
+ Get the state of the associated link.
+
+ Args:
+ body_id (int): body unique id.
+ link_id (int): link index.
+ compute_velocity (bool): If True, the Cartesian world velocity will be computed and returned.
+ compute_forward_kinematics (bool): if True, the Cartesian world position/orientation will be recomputed
+ using forward kinematics.
+
+ Returns:
+ np.float[3]: Cartesian position of CoM
+ np.float[4]: Cartesian orientation of CoM, in quaternion [x,y,z,w]
+ np.float[3]: local position offset of inertial frame (center of mass) expressed in the URDF link frame
+ np.float[4]: local orientation (quaternion [x,y,z,w]) offset of the inertial frame expressed in URDF link
+ frame
+ np.float[3]: world position of the URDF link frame
+ np.float[4]: world orientation of the URDF link frame
+ np.float[3]: Cartesian world linear velocity. Only returned if `compute_velocity` is True.
+ np.float[3]: Cartesian world angular velocity. Only returned if `compute_velocity` is True.
+ """
+ results = self.sim.getLinkState(body_id, link_id, computeLinkVelocity=compute_velocity,
+ computeForwardKinematics=compute_forward_kinematics)
+ return [np.array(result) for result in results]
+
+ def get_link_states(self, body_id, link_ids, compute_velocity=False, compute_forward_kinematics=False):
+ """
+ Get the state of the associated links.
+
+ Args:
+ body_id (int): body unique id.
+ link_ids (list of int): list of link index.
+ compute_velocity (bool): If True, the Cartesian world velocity will be computed and returned.
+ compute_forward_kinematics (bool): if True, the Cartesian world position/orientation will be recomputed
+ using forward kinematics.
+
+ Returns:
+ list:
+ np.float[3]: Cartesian position of CoM
+ np.float[4]: Cartesian orientation of CoM, in quaternion [x,y,z,w]
+ np.float[3]: local position offset of inertial frame (center of mass) expressed in the URDF link frame
+ np.float[4]: local orientation (quaternion [x,y,z,w]) offset of the inertial frame expressed in URDF
+ link frame
+ np.float[3]: world position of the URDF link frame
+ np.float[4]: world orientation of the URDF link frame
+ np.float[3]: Cartesian world linear velocity. Only returned if `compute_velocity` is True.
+ np.float[3]: Cartesian world angular velocity. Only returned if `compute_velocity` is True.
+ """
+ return [self.get_link_state(body_id, link_id, compute_velocity, compute_forward_kinematics)
+ for link_id in link_ids]
+
+ def get_link_names(self, body_id, link_ids):
+ """
+ Return the name of the given link(s).
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (int, list of int): link id, or list of link ids.
+
+ Returns:
+ if 1 link:
+ str: link name
+ if multiple links:
+ str[N]: link names
+ """
+ if isinstance(link_ids, int):
+ return self.sim.getJointInfo(body_id, link_ids)[12]
+ return [self.sim.getJointInfo(body_id, link_id)[12] for link_id in link_ids]
+
+ def get_link_masses(self, body_id, link_ids):
+ """
+ Return the mass of the given link(s).
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (int, list of int): link id, or list of link ids.
+
+ Returns:
+ if 1 link:
+ float: mass of the given link
+ else:
+ float[N]: mass of each link
+ """
+ if isinstance(link_ids, int):
+ return self.sim.getDynamicsInfo(body_id, link_ids)[0]
+ return np.array([self.sim.getDynamicsInfo(body_id, link_id)[0] for link_id in link_ids])
+
+ def get_link_frames(self, body_id, link_ids):
+ pass
+
+ def get_link_world_positions(self, body_id, link_ids):
+ """
+ Return the CoM position (in the Cartesian world space coordinates) of the given link(s).
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (list of int): list of link indices.
+
+ Returns:
+ if 1 link:
+ np.float[3]: the link CoM position in the world space
+ if multiple links:
+ np.float[N,3]: CoM position of each link in world space
+ """
+ if isinstance(link_ids, int):
+ if link_ids == -1:
+ return self.get_base_position(body_id)
+ return np.array(self.sim.getLinkState(body_id, link_ids)[0])
+ positions = []
+ for link_id in link_ids:
+ if link_id == -1:
+ positions.append(self.get_base_position(body_id))
+ else:
+ positions.append(np.array(self.sim.getLinkState(body_id, link_id)[0]))
+ return np.array(positions)
+
+ def get_link_positions(self, body_id, link_ids):
+ pass
+
+ def get_link_world_orientations(self, body_id, link_ids):
+ """
+ Return the CoM orientation (in the Cartesian world space) of the given link(s).
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (list of int): list of link indices.
+
+ Returns:
+ if 1 link:
+ np.float[4]: Cartesian orientation of the link CoM (x,y,z,w)
+ if multiple links:
+ np.float[N,4]: CoM orientation of each link (x,y,z,w)
+ """
+ if isinstance(link_ids, int):
+ if link_ids == -1:
+ return self.get_base_orientation(body_id)
+ return np.array(self.sim.getLinkState(body_id, link_ids)[1])
+ orientations = []
+ for link_id in link_ids:
+ if link_id == -1:
+ orientations.append(self.get_base_orientation(body_id))
+ else:
+ orientations.append(np.array(self.sim.getLinkState(body_id, link_id)[1]))
+ return np.array(orientations)
+
+ def get_link_orientations(self, body_id, link_ids):
+ pass
+
+ def get_link_world_linear_velocities(self, body_id, link_ids):
+ """
+ Return the linear velocity of the link(s) expressed in the Cartesian world space coordinates.
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (list of int): list of link indices.
+
+ Returns:
+ if 1 link:
+ np.float[3]: linear velocity of the link in the Cartesian world space
+ if multiple links:
+ np.float[N,3]: linear velocity of each link
+ """
+ if isinstance(link_ids, int):
+ if link_ids == -1:
+ return self.get_base_linear_velocity(body_id)
+ return np.array(self.sim.getLinkState(body_id, link_ids, computeLinkVelocity=1)[6])
+ velocities = []
+ for link_id in link_ids:
+ if link_id == -1:
+ velocities.append(self.get_base_linear_velocity(body_id))
+ else:
+ velocities.append(np.array(self.sim.getLinkState(body_id, link_id, computeLinkVelocity=1)[6]))
+ return np.array(velocities)
+
+ def get_link_world_angular_velocities(self, body_id, link_ids):
+ """
+ Return the angular velocity of the link(s) in the Cartesian world space coordinates.
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (list of int): list of link indices.
+
+ Returns:
+ if 1 link:
+ np.float[3]: angular velocity of the link in the Cartesian world space
+ if multiple links:
+ np.float[N,3]: angular velocity of each link
+ """
+ if isinstance(link_ids, int):
+ if link_ids == -1:
+ return self.get_base_linear_velocity(body_id)
+ return np.array(self.sim.getLinkState(body_id, link_ids, computeLinkVelocity=1)[7])
+ velocities = []
+ for link_id in link_ids:
+ if link_id == -1:
+ velocities.append(self.get_base_linear_velocity(body_id))
+ else:
+ velocities.append(np.array(self.sim.getLinkState(body_id, link_id, computeLinkVelocity=1)[7]))
+ return np.array(velocities)
+
+ def get_link_world_velocities(self, body_id, link_ids):
+ """
+ Return the linear and angular velocities (expressed in the Cartesian world space coordinates) for the given
+ link(s).
+
+ Args:
+ body_id (int): unique body id.
+ link_ids (list of int): list of link indices.
+
+ Returns:
+ if 1 link:
+ np.float[6]: linear and angular velocity of the link in the Cartesian world space
+ if multiple links:
+ np.float[N,6]: linear and angular velocity of each link
+ """
+ if isinstance(link_ids, int):
+ if link_ids == -1:
+ lin_vel, ang_vel = self.get_base_velocity(body_id)
+ return np.concatenate((lin_vel, ang_vel))
+ lin_vel, ang_vel = self.sim.getLinkState(body_id, link_ids, computeLinkVelocity=1)[6:8]
+ return np.array(lin_vel + ang_vel)
+ velocities = []
+ for link_id in link_ids:
+ if link_id == -1: # base link
+ lin_vel, ang_vel = self.get_base_velocity(body_id)
+ else:
+ lin_vel, ang_vel = self.sim.getLinkState(body_id, link_id, computeLinkVelocity=1)[6:8]
+ velocities.append(np.concatenate((lin_vel, ang_vel)))
+ return np.array(velocities)
+
+ def get_link_velocities(self, body_id, link_ids):
+ pass
+
+ def get_qindex(self, body_id, joint_ids):
+ """
+ Get the corresponding q index of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ int: q index
+ if multiple joints:
+ np.int[N]: q indices
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointInfo(body_id, joint_ids)[3] - 7
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[3] for joint_id in joint_ids]) - 7
+
+ def get_actuated_joint_ids(self, body_id):
+ """
+ Get the actuated joint ids associated with the given body id.
+
+ Warnings: this checks through the list of all joints each time it is called. It might be a good idea to call
+ this method one time and cache the actuated joint ids.
+
+ Args:
+ body_id (int): unique body id.
+
+ Returns:
+ list of int: actuated joint ids.
+ """
+ joint_ids = []
+ for joint_id in range(self.num_joints(body_id)):
+ # Get joint info
+ jnt = self.get_joint_info(body_id, joint_id)
+ if jnt[2] != self.sim.JOINT_FIXED: # if not a fixed joint
+ joint_ids.append(jnt[0])
+ return joint_ids
+
+ def get_joint_names(self, body_id, joint_ids):
+ """
+ Return the name of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ str: name of the joint
+ if multiple joints:
+ str[N]: name of each joint
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointInfo(body_id, joint_ids)[1]
+ return [self.sim.getJointInfo(body_id, joint_id)[1] for joint_id in joint_ids]
+
+ def get_joint_dampings(self, body_id, joint_ids):
+ """
+ Get the damping coefficient of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: damping coefficient of the given joint
+ if multiple joints:
+ np.float[N]: damping coefficient for each specified joint
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointInfo(body_id, joint_ids)[6]
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[6] for joint_id in joint_ids])
+
+ def get_joint_frictions(self, body_id, joint_ids):
+ """
+ Get the friction coefficient of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: friction coefficient of the given joint
+ if multiple joints:
+ float[N]: friction coefficient for each specified joint
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointInfo(body_id, joint_ids)[7]
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[7] for joint_id in joint_ids])
+
+ def get_joint_limits(self, body_id, joint_ids):
+ """
+ Get the joint limits of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ np.float[2]: lower and upper limit
+ if multiple joints:
+ np.float[N,2]: lower and upper limit for each specified joint
+ """
+ if isinstance(joint_ids, int):
+ return np.array(self.sim.getJointInfo(body_id, joint_ids)[8:10])
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[8:10] for joint_id in joint_ids])
+
+ def get_joint_max_forces(self, body_id, joint_ids):
+ """
+ Get the maximum force that can be applied on the given joint(s).
+
+ Warning: Note that this is not automatically used in position, velocity, or torque control.
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: maximum force [N]
+ if multiple joints:
+ float[N]: maximum force for each specified joint [N]
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointInfo(body_id, joint_ids)[10]
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[10] for joint_id in joint_ids])
+
+ def get_joint_max_velocities(self, body_id, joint_ids):
+ """
+ Get the maximum velocity that can be applied on the given joint(s).
+
+ Warning: Note that this is not automatically used in position, velocity, or torque control.
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: maximum velocity [rad/s]
+ if multiple joints:
+ np.float[N]: maximum velocities for each specified joint [rad/s]
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointInfo(body_id, joint_ids)[11]
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[11] for joint_id in joint_ids])
+
+ def get_joint_axes(self, body_id, joint_ids):
+ """
+ Get the joint axis about the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ np.float[3]: joint axis
+ if multiple joint:
+ np.float[N,3]: list of joint axis
+ """
+ if isinstance(joint_ids, int):
+ return np.array(self.sim.getJointInfo(body_id, joint_ids)[-4])
+ return np.array([self.sim.getJointInfo(body_id, joint_id)[-4] for joint_id in joint_ids])
+
+ def set_joint_positions(self, body_id, joint_ids, positions, velocities=None, kps=None, kds=None, forces=None):
+ """
+ Set the position of the given joint(s) (using position control).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+ positions (float, np.float[N]): desired position, or list of desired positions [rad]
+ velocities (None, float, np.float[N]): desired velocity, or list of desired velocities [rad/s]
+ kps (None, float, np.float[N]): position gain(s)
+ kds (None, float, np.float[N]): velocity gain(s)
+ forces (float): maximum motor force(s)/torque(s) used to reach the target values.
+ """
+ if isinstance(joint_ids, int):
+ self.set_joint_motor_control(body_id, joint_ids, control_mode=pybullet.POSITION_CONTROL, position=positions,
+ velocity=velocities, force=forces, kp=kps, kd=kds)
+ else:
+ self.set_joint_motor_control_array(body_id, joint_ids, control_mode=pybullet.POSITION_CONTROL,
+ positions=positions, velocities=velocities, forces=forces, kps=kps,
+ kds=kds)
+
+ def get_joint_positions(self, body_id, joint_ids):
+ """
+ Get the position of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: joint position [rad]
+ if multiple joints:
+ np.float[N]: joint positions [rad]
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointState(body_id, joint_ids)[0]
+ return np.array([state[0] for state in self.sim.getJointStates(body_id, joint_ids)])
+
+ def set_joint_velocities(self, body_id, joint_ids, velocities, max_force=None):
+ """
+ Set the velocity of the given joint(s) (using velocity control).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+ velocities (float, np.float[N]): desired velocity, or list of desired velocities [rad/s]
+ max_force (bool, float, float[N]): maximum motor forces/torques
+ """
+ if isinstance(joint_ids, int):
+ if max_force is None:
+ self.sim.setJointMotorControl2(body_id, joint_ids, self.sim.VELOCITY_CONTROL, targetVelocity=velocities)
+ self.sim.setJointMotorControl2(body_id, joint_ids, self.sim.VELOCITY_CONTROL, targetVelocity=velocities,
+ force=max_force)
+ if max_force is None:
+ self.sim.setJointMotorControlArray(body_id, joint_ids, self.sim.VELOCITY_CONTROL,
+ targetVelocities=velocities)
+ self.sim.setJointMotorControlArray(body_id, joint_ids, self.sim.VELOCITY_CONTROL,
+ targetVelocities=velocities, forces=max_force)
+
+ def get_joint_velocities(self, body_id, joint_ids):
+ """
+ Get the velocity of the given joint(s).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: joint velocity [rad/s]
+ if multiple joints:
+ np.float[N]: joint velocities [rad/s]
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointState(body_id, joint_ids)[1]
+ return np.array([state[1] for state in self.sim.getJointStates(body_id, joint_ids)])
+
+ def set_joint_accelerations(self, body_id, joint_ids, accelerations, q=None, dq=None):
+ """
+ Set the acceleration of the given joint(s) (using force control). This is achieved by performing inverse
+ dynamic which given the joint accelerations compute the joint torques to be applied.
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+ accelerations (float, np.float[N]): desired joint acceleration, or list of desired joint accelerations
+ [rad/s^2]
+ """
+ # check joint ids
+ if isinstance(joint_ids, int):
+ joint_ids = [joint_ids]
+ if isinstance(accelerations, (int, float)):
+ accelerations = [accelerations]
+ if len(accelerations) != len(joint_ids):
+ raise ValueError("Expecting the desired accelerations to be of the same size as the number of joints; "
+ "{} != {}".format(len(accelerations), len(joint_ids)))
+
+ # get position and velocities
+ if q is None or dq is None:
+ joints = self.get_actuated_joint_ids(body_id)
+ if q is None:
+ q = self.get_joint_positions(body_id, joints)
+ if dq is None:
+ dq = self.get_joint_velocities(body_id, joints)
+
+ num_actuated_joints = len(q)
+
+ # if joint accelerations vector is not the same size as the actuated joints
+ if len(accelerations) != num_actuated_joints:
+ q_idx = self.get_qindex(joint_ids)
+ acc = np.zeros(num_actuated_joints)
+ acc[q_idx] = accelerations
+ accelerations = acc
+
+ # compute joint torques from Inverse Dynamics
+ torques = self.calculate_inverse_dynamics(body_id, q, dq, accelerations)
+
+ # get corresponding torques
+ if len(torques) != len(joint_ids):
+ q_idx = self.get_qindex(joint_ids)
+ torques = torques[q_idx]
+
+ # set the joint torques
+ self.set_joint_torques(body_id, joint_ids, torques)
+
+ def get_joint_accelerations(self, body_id, joint_ids, q=None, dq=None):
+ """
+ Get the acceleration at the given joint(s). This is carried out by first getting the joint torques, then
+ performing forward dynamics to get the joint accelerations from the joint torques.
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+ q (list of int, None): all the joint positions. If None, it will compute it.
+ dq (list of int, None): all the joint velocities. If None, it will compute it.
+
+ Returns:
+ if 1 joint:
+ float: joint acceleration [rad/s^2]
+ if multiple joints:
+ np.float[N]: joint accelerations [rad/s^2]
+ """
+ # get the torques
+ torques = self.get_joint_torques(body_id, joint_ids)
+
+ # get position and velocities
+ if q is None or dq is None:
+ joints = self.get_actuated_joint_ids(body_id)
+ if q is None:
+ q = self.get_joint_positions(body_id, joints)
+ if dq is None:
+ dq = self.get_joint_velocities(body_id, joints)
+
+ # compute the accelerations
+ accelerations = self.calculate_forward_dynamics(body_id, q, dq, torques=torques)
+
+ # return the specified accelerations
+ q_idx = self.get_qindex(body_id, joint_ids)
+ return accelerations[q_idx]
+
+ def set_joint_torques(self, body_id, joint_ids, torques):
+ """
+ Set the torque/force to the given joint(s) (using force/torque control).
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): joint id, or list of joint ids.
+ torque (float, list of float): desired torque(s) to apply to the joint(s) [N].
+ """
+ if isinstance(joint_ids, int):
+ self.sim.setJointMotorControl2(body_id, joint_ids, self.sim.TORQUE_CONTROL, force=torques)
+ self.sim.setJointMotorControlArray(body_id, joint_ids, self.sim.TORQUE_CONTROL, forces=torques)
+
+ def get_joint_torques(self, body_id, joint_ids):
+ """
+ Get the applied torque(s) on the given joint(s). "This is the motor torque applied during the last `step`.
+ Note that this only applies in VELOCITY_CONTROL and POSITION_CONTROL. If you use TORQUE_CONTROL then the
+ applied joint motor torque is exactly what you provide, so there is no need to report it separately." [1]
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, list of int): a joint id, or list of joint ids.
+
+ Returns:
+ if 1 joint:
+ float: torque [Nm]
+ if multiple joints:
+ np.float[N]: torques associated to the given joints [Nm]
+ """
+ if isinstance(joint_ids, int):
+ return self.sim.getJointState(body_id, joint_ids)[3]
+ return np.array([state[3] for state in self.sim.getJointStates(body_id, joint_ids)])
+
+ def get_joint_reaction_forces(self, body_id, joint_ids):
+ """
+ Return the joint reaction forces at the given joint. Note that the torque sensor must be enabled, otherwise
+ it will always return [0,0,0,0,0,0].
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, int[N]): joint id, or list of joint ids
+
+ Returns:
+ if 1 joint:
+ np.float[6]: joint reaction force (fx,fy,fz,mx,my,mz) [N,Nm]
+ if multiple joints:
+ np.float[N,6]: joint reaction forces [N, Nm]
+ """
+ if isinstance(joint_ids, int):
+ return np.array(self.sim.getJointState(body_id, joint_ids)[2])
+ return np.array([state[2] for state in self.sim.getJointStates(body_id, joint_ids)])
+
+ def get_joint_powers(self, body_id, joint_ids):
+ """
+ Return the applied power at the given joint(s). Power = torque * velocity.
+
+ Args:
+ body_id (int): unique body id.
+ joint_ids (int, int[N]): joint id, or list of joint ids
+
+ Returns:
+ if 1 joint:
+ float: joint power [W]
+ if multiple joints:
+ np.float[N]: power at each joint [W]
+ """
+ torque = self.get_joint_torques(body_id, joint_ids)
+ velocity = self.get_joint_velocities(body_id, joint_ids)
+ return torque * velocity
+
+ #################
+ # visualization #
+ #################
+
+ def create_visual_shape(self, shape_type, radius=0.5, half_extents=(1., 1., 1.), length=1., filename=None,
+ mesh_scale=(1., 1., 1.), plane_normal=(0., 0., 1.), flags=-1, rgba_color=None,
+ specular_color=None, visual_frame_position=None, vertices=None, indices=None, uvs=None,
+ normals=None, visual_frame_orientation=None):
+ """
+ Create a visual shape in the simulator.
+
+ Args:
+ shape_type (int): type of shape; GEOM_SPHERE (=2), GEOM_BOX (=3), GEOM_CAPSULE (=7), GEOM_CYLINDER (=4),
+ GEOM_PLANE (=6), GEOM_MESH (=5)
+ radius (float): only for GEOM_SPHERE, GEOM_CAPSULE, GEOM_CYLINDER
+ half_extents (np.float[3], list/tuple of 3 floats): only for GEOM_BOX.
+ length (float): only for GEOM_CAPSULE, GEOM_CYLINDER (length = height).
+ filename (str): Filename for GEOM_MESH, currently only Wavefront .obj. Will create convex hulls for each
+ object (marked as 'o') in the .obj file.
+ mesh_scale (np.float[3], list/tuple of 3 floats): scale of mesh (only for GEOM_MESH).
+ plane_normal (np.float[3], list/tuple of 3 floats): plane normal (only for GEOM_PLANE).
+ flags (int): unused / to be decided
+ rgba_color (list/tuple of 4 floats): color components for red, green, blue and alpha, each in range [0..1].
+ specular_color (list/tuple of 3 floats): specular reflection color, red, green, blue components in range
+ [0..1]
+ visual_frame_position (np.float[3]): translational offset of the visual shape with respect to the link frame
+ vertices (list of np.float[3]): Instead of creating a mesh from obj file, you can provide vertices, indices,
+ uvs and normals
+ indices (list of int): triangle indices, should be a multiple of 3.
+ uvs (list of np.float[2]): uv texture coordinates for vertices. Use changeVisualShape to choose the
+ texture image. The number of uvs should be equal to number of vertices
+ normals (list of np.float[3]): vertex normals, number should be equal to number of vertices.
+ visual_frame_orientation (np.float[4]): rotational offset (quaternion x,y,z,w) of the visual shape with
+ respect to the link frame
+
+ Returns:
+ int: The return value is a non-negative int unique id for the visual shape or -1 if the call failed.
+ """
+ # add few variables
+ kwargs = {}
+ if rgba_color is not None:
+ kwargs['rgbaColor'] = rgba_color
+ if specular_color is not None:
+ kwargs['specularColor'] = specular_color
+ if visual_frame_position is not None:
+ kwargs['visualFramePosition'] = visual_frame_position
+ if visual_frame_orientation is not None:
+ kwargs['visualFrameOrientation'] = visual_frame_orientation
+
+ if shape_type == self.sim.GEOM_SPHERE:
+ return self.sim.createVisualShape(shape_type, radius=radius, **kwargs)
+ elif shape_type == self.sim.GEOM_BOX:
+ return self.sim.createVisualShape(shape_type, halfExtents=half_extents, **kwargs)
+ elif shape_type == self.sim.GEOM_CAPSULE or shape_type == self.sim.GEOM_CYLINDER:
+ return self.sim.createVisualShape(shape_type, radius=radius, length=length, **kwargs)
+ elif shape_type == self.sim.GEOM_PLANE:
+ return self.sim.createVisualShape(shape_type, planeNormal=plane_normal, **kwargs)
+ elif shape_type == self.sim.GEOM_MESH:
+ if filename is not None:
+ kwargs['fileName'] = filename
+ else:
+ if vertices is not None:
+ kwargs['vertices'] = vertices
+ if indices is not None:
+ kwargs['indices'] = indices
+ if uvs is not None:
+ kwargs['uvs'] = uvs
+ if normals is not None:
+ kwargs['normals'] = normals
+ return self.sim.createVisualShape(shape_type, **kwargs)
+ else:
+ raise ValueError("Unknown visual shape type.")
+
+ def get_visual_shape_data(self, object_id, flags=-1):
+ """
+ Get the visual shape data associated with the given object id.
+
+ Args:
+ object_id (int): object unique id.
+ flags (int, None): VISUAL_SHAPE_DATA_TEXTURE_UNIQUE_IDS (=1) will also provide `texture_unique_id`.
+
+ Returns:
+ int: object unique id.
+ int: link index or -1 for the base
+ int: visual geometry type (TBD)
+ np.float[3]: dimensions (size, local scale) of the geometry
+ str: path to the triangle mesh, if any. Typically relative to the URDF, SDF or MJCF file location, but
+ could be absolute
+ np.float[3]: position of local visual frame, relative to link/joint frame
+ np.float[4]: orientation of local visual frame relative to link/joint frame
+ list of 4 floats: URDF color (if any specified) in Red / Green / Blue / Alpha
+ int: texture unique id of the shape or -1 if None. This field only exists if using
+ VISUAL_SHAPE_DATA_TEXTURE_UNIQUE_IDS (=1) flag.
+ """
+ return self.sim.getVisualShapeData(object_id, flags=flags)
+
+ def change_visual_shape(self, object_id, link_id, shape_id=None, texture_id=None, rgba_color=None,
+ specular_color=None):
+ """
+ Allows to change the texture of a shape, the RGBA color and other properties.
+
+ Args:
+ object_id (int): unique object id.
+ link_id (int): link id.
+ shape_id (int): shape id.
+ texture_id (int): texture id.
+ rgba_color (float[4]): RGBA color. Each is in the range [0..1]. Alpha has to be 0 (invisible) or 1
+ (visible) at the moment.
+ specular_color (int[3]): specular color components, RED, GREEN and BLUE, can be from 0 to large number
+ (>100).
+ """
+ kwargs = {}
+ if shape_id is not None:
+ kwargs['shapeIndex'] = shape_id
+ if texture_id is not None:
+ kwargs['textureUniqueId'] = texture_id
+ if rgba_color is not None:
+ kwargs['rgbaColor'] = rgba_color
+ if specular_color is not None:
+ kwargs['specularColor'] = specular_color
+ self.sim.changeVisualShape(object_id, link_id, **kwargs)
+
+ def load_texture(self, filename):
+ """
+ Load a texture from file and return a non-negative texture unique id if the loading succeeds.
+ This unique id can be used with changeVisualShape.
+
+ Args:
+ filename (str): path to the file.
+
+ Returns:
+ int: texture unique id. If non-negative, the texture was loaded successfully.
+ """
+ return self.sim.loadTexture(filename)
+
+ def compute_view_matrix(self, eye_position, target_position, up_vector):
+ """Compute the view matrix.
+
+ The view matrix is the 4x4 matrix that maps the world coordinates into the camera coordinates. Basically,
+ it applies a rotation and translation such that the world is in front of the camera. That is, instead
+ of turning the camera to capture what we want in the world, we keep the camera fixed and turn the world.
+
+ Args:
+ eye_position (np.float[3]): eye position in Cartesian world coordinates
+ target_position (np.float[3]): position of the target (focus) point in Cartesian world coordinates
+ up_vector (np.float[3]): up vector of the camera in Cartesian world coordinates
+
+ Returns:
+ np.float[4,4]: the view matrix
+
+ More info:
+ [1] http://www.codinglabs.net/article_world_view_projection_matrix.aspx
+ [2] http://www.thecodecrate.com/opengl-es/opengl-transformation-matrices/
+ """
+ view = self.sim.computeViewMatrix(eyePosition=eye_position, targetPosition=target_position, upVector=up_vector)
+ return np.array(view).reshape(4, 4).T
+
+ def compute_view_matrix_from_ypr(self, target_position, distance, yaw, pitch, roll, up_axis_index=2):
+ """Compute the view matrix from the yaw, pitch, and roll angles.
+
+ The view matrix is the 4x4 matrix that maps the world coordinates into the camera coordinates. Basically,
+ it applies a rotation and translation such that the world is in front of the camera. That is, instead
+ of turning the camera to capture what we want in the world, we keep the camera fixed and turn the world.
+
+ Args:
+ target_position (np.float[3]): target focus point in Cartesian world coordinates
+ distance (float): distance from eye to focus point
+ yaw (float): yaw angle in radians left/right around up-axis
+ pitch (float): pitch in radians up/down.
+ roll (float): roll in radians around forward vector
+ up_axis_index (int): either 1 for Y or 2 for Z axis up.
+
+ Returns:
+ np.float[4,4]: the view matrix
+
+ More info:
+ [1] http://www.codinglabs.net/article_world_view_projection_matrix.aspx
+ [2] http://www.thecodecrate.com/opengl-es/opengl-transformation-matrices/
+ """
+ view = self.sim.computeViewMatrixFromYawPitchRoll(targetPosition=target_position, distance=distance,
+ yaw=np.rad2deg(yaw), pitch=np.rad2deg(pitch),
+ roll=np.rad2deg(roll), upAxisIndex=up_axis_index)
+ return np.array(view).reshape(4, 4).T
+
+ def compute_projection_matrix(self, left, right, bottom, top, near, far):
+ """Compute the orthographic projection matrix.
+
+ The projection matrix is the 4x4 matrix that maps from the camera/eye coordinates to clipped coordinates.
+ It is applied after the view matrix.
+
+ There are 2 projection matrices:
+ * orthographic projection
+ * perspective projection
+
+ For the perspective projection, see `computeProjectionMatrixFOV(self)
+
+ Args:
+ left (float): left screen (canvas) coordinate
+ right (float): right screen (canvas) coordinate
+ bottom (float): bottom screen (canvas) coordinate
+ top (float): top screen (canvas) coordinate
+ near (float): near plane distance
+ far (float): far plane distance
+
+ Returns:
+ np.float[4,4]: the perspective projection matrix
+
+ More info:
+ [1] http://www.codinglabs.net/article_world_view_projection_matrix.aspx
+ [2] http://www.thecodecrate.com/opengl-es/opengl-transformation-matrices/
+ """
+ proj = self.sim.computeProjectionMatrix(left, right, bottom, top, near, far)
+ return np.array(proj).reshape(4, 4).T
+
+ def compute_projection_matrix_fov(self, fov, aspect, near, far):
+ """Compute the perspective projection matrix using the field of view (FOV).
+
+ Args:
+ fov (float): field of view
+ aspect (float): aspect ratio
+ near (float): near plane distance
+ far (float): far plane distance
+
+ Returns:
+ np.float[4,4]: the perspective projection matrix
+
+ More info:
+ [1] http://www.codinglabs.net/article_world_view_projection_matrix.aspx
+ [2] http://www.thecodecrate.com/opengl-es/opengl-transformation-matrices/
+ """
+ proj = self.sim.computeProjectionMatrixFOV(fov, aspect, near, far)
+ return np.array(proj).reshape(4, 4).T
+
+ def get_camera_image(self, width, height, view_matrix=None, projection_matrix=None, light_direction=None,
+ light_color=None, light_distance=None, shadow=None, light_ambient_coeff=None,
+ light_diffuse_coeff=None, light_specular_coeff=None, renderer=None, flags=None):
+ """
+ The `get_camera_image` API will return a RGB image, a depth buffer and a segmentation mask buffer with body
+ unique ids of visible objects for each pixel. Note that PyBullet can be compiled using the numpy option:
+ using numpy will improve the performance of copying the camera pixels from C to Python.
+
+ Note that copying pixels from C/C++ to Python can be really slow for large images, unless you compile PyBullet
+ using NumPy. You can check if NumPy is enabled using `PyBullet.isNumpyEnabled()`. `pip install pybullet` has
+ NumPy enabled, if available on the system.
+
+ Args:
+ width (int): horizontal image resolution in pixels
+ height (int): vertical image resolution in pixels
+ view_matrix (np.float[4,4]): 4x4 view matrix, see `compute_view_matrix`
+ projection_matrix (np.float[4,4]): 4x4 projection matrix, see `compute_projection`
+ light_direction (np.float[3]): `light_direction` specifies the world position of the light source,
+ the direction is from the light source position to the origin of the world frame.
+ light_color (np.float[3]): directional light color in [RED,GREEN,BLUE] in range 0..1
+ light_distance (float): distance of the light along the normalized `light_direction`
+ shadow (bool): True for shadows, False for no shadows
+ light_ambient_coeff (float): light ambient coefficient
+ light_diffuse_coeff (float): light diffuse coefficient
+ light_specular_coeff (float): light specular coefficient
+ renderer (int): ER_BULLET_HARDWARE_OPENGL (=131072) or ER_TINY_RENDERER (=65536). Note that DIRECT (=2)
+ mode has no OpenGL, so it requires ER_TINY_RENDERER (=65536).
+ flags (int): ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX (=1), See below in description of
+ segmentationMaskBuffer and example code. Use ER_NO_SEGMENTATION_MASK (=4) to avoid calculating the
+ segmentation mask.
+
+ Returns:
+ int: width image resolution in pixels (horizontal)
+ int: height image resolution in pixels (vertical)
+ np.int[width, height, 4]: RBGA pixels (each pixel is in the range [0..255] for each channel R, G, B, A)
+ np.float[width, heigth]: Depth buffer. Bullet uses OpenGL to render, and the convention is non-linear
+ z-buffer. See https://stackoverflow.com/questions/6652253/getting-the-true-z-value-from-the-depth-buffer
+ Using the projection matrix, the depth is computed as:
+ `depth = far * near / (far - (far - near) * depthImg)`, where `depthImg` is the depth from Bullet
+ `getCameraImage`, far=1000. and near=0.01.
+ np.int[width, height]: Segmentation mask buffer. For each pixels the visible object unique id.
+ If ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX (=1) is used, the segmentationMaskBuffer combines the
+ object unique id and link index as follows: value = objectUniqueId + (linkIndex+1)<<24.
+ So for a free floating body without joints/links, the segmentation mask is equal to its body unique id,
+ since its link index is -1.
+ """
+ kwargs = {}
+ if view_matrix is not None:
+ if isinstance(view_matrix, np.ndarray):
+ kwargs['viewMatrix'] = view_matrix.T.ravel().tolist()
+ else:
+ kwargs['viewMatrix'] = view_matrix
+ if projection_matrix is not None:
+ if isinstance(projection_matrix, np.ndarray):
+ kwargs['projectionMatrix'] = projection_matrix.T.ravel().tolist()
+ else:
+ kwargs['projectionMatrix'] = projection_matrix
+ if light_direction is not None:
+ if isinstance(light_direction, np.ndarray):
+ kwargs['lightDirection'] = light_direction.ravel().tolist()
+ else:
+ kwargs['lightDirection'] = light_direction
+ if light_color is not None:
+ if isinstance(light_color, np.ndarray):
+ kwargs['lightColor'] = light_color
+ else:
+ kwargs['lightColor'] = light_color
+ if light_distance is not None:
+ kwargs['lightDistance'] = light_distance
+ if shadow is not None:
+ kwargs['shadow'] = int(shadow)
+ if light_ambient_coeff is not None:
+ kwargs['lightAmbientCoeff'] = light_ambient_coeff
+ if light_diffuse_coeff is not None:
+ kwargs['lightDiffuseCoeff'] = light_diffuse_coeff
+ if light_specular_coeff is not None:
+ kwargs['lightSpecularCoeff'] = light_specular_coeff
+ if renderer is not None:
+ kwargs['renderer'] = renderer
+ if flags is not None:
+ kwargs['flags'] = flags
+
+ width, height, rgba, depth, segmentation = self.sim.getCameraImage(width, height, **kwargs)
+ rgba = np.array(rgba).reshape(width, height, 4)
+ depth = np.array(depth).reshape(width, height)
+ segmentation = np.array(segmentation).reshape(width, height)
+ return width, height, rgba, depth, segmentation
+
+ def get_rgba_image(self, width, height, view_matrix=None, projection_matrix=None, light_direction=None,
+ light_color=None, light_distance=None, shadow=None, light_ambient_coeff=None,
+ light_diffuse_coeff=None, light_specular_coeff=None, renderer=None, flags=None):
+ """
+ The `get_rgba_image` API will return a RGBA image. Note that PyBullet can be compiled using the numpy option:
+ using numpy will improve the performance of copying the camera pixels from C to Python.
+
+ Note that copying pixels from C/C++ to Python can be really slow for large images, unless you compile PyBullet
+ using NumPy. You can check if NumPy is enabled using `PyBullet.isNumpyEnabled()`. `pip install pybullet` has
+ NumPy enabled, if available on the system.
+
+ Args:
+ width (int): horizontal image resolution in pixels
+ height (int): vertical image resolution in pixels
+ view_matrix (np.float[4,4]): 4x4 view matrix, see `compute_view_matrix`
+ projection_matrix (np.float[4,4]): 4x4 projection matrix, see `compute_projection`
+ light_direction (np.float[3]): `light_direction` specifies the world position of the light source,
+ the direction is from the light source position to the origin of the world frame.
+ light_color (np.float[3]): directional light color in [RED,GREEN,BLUE] in range 0..1
+ light_distance (float): distance of the light along the normalized `light_direction`
+ shadow (bool): True for shadows, False for no shadows
+ light_ambient_coeff (float): light ambient coefficient
+ light_diffuse_coeff (float): light diffuse coefficient
+ light_specular_coeff (float): light specular coefficient
+ renderer (int): ER_BULLET_HARDWARE_OPENGL (=131072) or ER_TINY_RENDERER (=65536). Note that DIRECT (=2)
+ mode has no OpenGL, so it requires ER_TINY_RENDERER (=65536).
+ flags (int): ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX (=1), See below in description of
+ segmentationMaskBuffer and example code. Use ER_NO_SEGMENTATION_MASK (=4) to avoid calculating the
+ segmentation mask.
+
+ Returns:
+ np.int[width, height, 4]: RBGA pixels (each pixel is in the range [0..255] for each channel R, G, B, A)
+ """
+ kwargs = {}
+ if view_matrix is not None:
+ if isinstance(view_matrix, np.ndarray):
+ kwargs['viewMatrix'] = view_matrix.T.ravel().tolist()
+ else:
+ kwargs['viewMatrix'] = view_matrix
+ if projection_matrix is not None:
+ if isinstance(projection_matrix, np.ndarray):
+ kwargs['projectionMatrix'] = projection_matrix.T.ravel().tolist()
+ else:
+ kwargs['projectionMatrix'] = projection_matrix
+ if light_direction is not None:
+ if isinstance(light_direction, np.ndarray):
+ kwargs['lightDirection'] = light_direction.ravel().tolist()
+ else:
+ kwargs['lightDirection'] = light_direction
+ if light_color is not None:
+ if isinstance(light_color, np.ndarray):
+ kwargs['lightColor'] = light_color
+ else:
+ kwargs['lightColor'] = light_color
+ if light_distance is not None:
+ kwargs['lightDistance'] = light_distance
+ if shadow is not None:
+ kwargs['shadow'] = int(shadow)
+ if light_ambient_coeff is not None:
+ kwargs['lightAmbientCoeff'] = light_ambient_coeff
+ if light_diffuse_coeff is not None:
+ kwargs['lightDiffuseCoeff'] = light_diffuse_coeff
+ if light_specular_coeff is not None:
+ kwargs['lightSpecularCoeff'] = light_specular_coeff
+ if renderer is not None:
+ kwargs['renderer'] = renderer
+ if flags is not None:
+ kwargs['flags'] = flags
+
+ img = np.array(self.sim.getCameraImage(width, height, **kwargs)[2])
+ img = img.reshape(width, height, 4) # RGBA
+ return img
+
+ def get_depth_image(self, width, height, view_matrix=None, projection_matrix=None, light_direction=None,
+ light_color=None, light_distance=None, shadow=None, light_ambient_coeff=None,
+ light_diffuse_coeff=None, light_specular_coeff=None, renderer=None, flags=None):
+ """
+ The `get_depth_image` API will return a depth buffer. Note that PyBullet can be compiled using the numpy option:
+ using numpy will improve the performance of copying the camera pixels from C to Python.
+
+ Note that copying pixels from C/C++ to Python can be really slow for large images, unless you compile PyBullet
+ using NumPy. You can check if NumPy is enabled using `PyBullet.isNumpyEnabled()`. `pip install pybullet` has
+ NumPy enabled, if available on the system.
+
+ Args:
+ width (int): horizontal image resolution in pixels
+ height (int): vertical image resolution in pixels
+ view_matrix (np.float[4,4]): 4x4 view matrix, see `compute_view_matrix`
+ projection_matrix (np.float[4,4]): 4x4 projection matrix, see `compute_projection`
+ light_direction (np.float[3]): `light_direction` specifies the world position of the light source,
+ the direction is from the light source position to the origin of the world frame.
+ light_color (np.float[3]): directional light color in [RED,GREEN,BLUE] in range 0..1
+ light_distance (float): distance of the light along the normalized `light_direction`
+ shadow (bool): True for shadows, False for no shadows
+ light_ambient_coeff (float): light ambient coefficient
+ light_diffuse_coeff (float): light diffuse coefficient
+ light_specular_coeff (float): light specular coefficient
+ renderer (int): ER_BULLET_HARDWARE_OPENGL (=131072) or ER_TINY_RENDERER (=65536). Note that DIRECT (=2)
+ mode has no OpenGL, so it requires ER_TINY_RENDERER (=65536).
+ flags (int): ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX (=1), See below in description of
+ segmentationMaskBuffer and example code. Use ER_NO_SEGMENTATION_MASK (=4) to avoid calculating the
+ segmentation mask.
+
+ Returns:
+ np.float[width, heigth]: Depth buffer. Bullet uses OpenGL to render, and the convention is non-linear
+ z-buffer. See https://stackoverflow.com/questions/6652253/getting-the-true-z-value-from-the-depth-buffer
+ Using the projection matrix, the depth is computed as:
+ `depth = far * near / (far - (far - near) * depthImg)`, where `depthImg` is the depth from Bullet
+ `getCameraImage`, far=1000. and near=0.01.
+ """
+ kwargs = {}
+ if view_matrix is not None:
+ if isinstance(view_matrix, np.ndarray):
+ kwargs['viewMatrix'] = view_matrix.T.ravel().tolist()
+ else:
+ kwargs['viewMatrix'] = view_matrix
+ if projection_matrix is not None:
+ if isinstance(projection_matrix, np.ndarray):
+ kwargs['projectionMatrix'] = projection_matrix.T.ravel().tolist()
+ else:
+ kwargs['projectionMatrix'] = projection_matrix
+ if light_direction is not None:
+ if isinstance(light_direction, np.ndarray):
+ kwargs['lightDirection'] = light_direction.ravel().tolist()
+ else:
+ kwargs['lightDirection'] = light_direction
+ if light_color is not None:
+ if isinstance(light_color, np.ndarray):
+ kwargs['lightColor'] = light_color
+ else:
+ kwargs['lightColor'] = light_color
+ if light_distance is not None:
+ kwargs['lightDistance'] = light_distance
+ if shadow is not None:
+ kwargs['shadow'] = int(shadow)
+ if light_ambient_coeff is not None:
+ kwargs['lightAmbientCoeff'] = light_ambient_coeff
+ if light_diffuse_coeff is not None:
+ kwargs['lightDiffuseCoeff'] = light_diffuse_coeff
+ if light_specular_coeff is not None:
+ kwargs['lightSpecularCoeff'] = light_specular_coeff
+ if renderer is not None:
+ kwargs['renderer'] = renderer
+ if flags is not None:
+ kwargs['flags'] = flags
+
+ img = np.array(self.sim.getCameraImage(width, height, **kwargs)[3])
+ img = img.reshape(width, height)
+ return img
+
+ def get_segmentation_image(self, width, height, view_matrix=None, projection_matrix=None, light_direction=None,
+ light_color=None, light_distance=None, shadow=None, light_ambient_coeff=None,
+ light_diffuse_coeff=None, light_specular_coeff=None, renderer=None, flags=None):
+ """
+ The `get_segmentation_image` API will return a segmentation mask buffer with body unique ids of visible objects
+ for each pixel. Note that PyBullet can be compiled using the numpy option: using numpy will improve
+ the performance of copying the camera pixels from C to Python.
+
+ Note that copying pixels from C/C++ to Python can be really slow for large images, unless you compile PyBullet
+ using NumPy. You can check if NumPy is enabled using `PyBullet.isNumpyEnabled()`. `pip install pybullet` has
+ NumPy enabled, if available on the system.
+
+ Args:
+ width (int): horizontal image resolution in pixels
+ height (int): vertical image resolution in pixels
+ view_matrix (np.float[4,4]): 4x4 view matrix, see `compute_view_matrix`
+ projection_matrix (np.float[4,4]): 4x4 projection matrix, see `compute_projection`
+ light_direction (np.float[3]): `light_direction` specifies the world position of the light source,
+ the direction is from the light source position to the origin of the world frame.
+ light_color (np.float[3]): directional light color in [RED,GREEN,BLUE] in range 0..1
+ light_distance (float): distance of the light along the normalized `light_direction`
+ shadow (bool): True for shadows, False for no shadows
+ light_ambient_coeff (float): light ambient coefficient
+ light_diffuse_coeff (float): light diffuse coefficient
+ light_specular_coeff (float): light specular coefficient
+ renderer (int): ER_BULLET_HARDWARE_OPENGL (=131072) or ER_TINY_RENDERER (=65536). Note that DIRECT (=2)
+ mode has no OpenGL, so it requires ER_TINY_RENDERER (=65536).
+ flags (int): ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX (=1), See below in description of
+ segmentationMaskBuffer and example code. Use ER_NO_SEGMENTATION_MASK (=4) to avoid calculating the
+ segmentation mask.
+
+ Returns:
+ np.int[width, height]: Segmentation mask buffer. For each pixels the visible object unique id.
+ If ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX (=1) is used, the segmentationMaskBuffer combines the
+ object unique id and link index as follows: value = objectUniqueId + (linkIndex+1)<<24.
+ So for a free floating body without joints/links, the segmentation mask is equal to its body unique id,
+ since its link index is -1.
+ """
+ kwargs = {}
+ if view_matrix is not None:
+ if isinstance(view_matrix, np.ndarray):
+ kwargs['viewMatrix'] = view_matrix.T.ravel().tolist()
+ else:
+ kwargs['viewMatrix'] = view_matrix
+ if projection_matrix is not None:
+ if isinstance(projection_matrix, np.ndarray):
+ kwargs['projectionMatrix'] = projection_matrix.T.ravel().tolist()
+ else:
+ kwargs['projectionMatrix'] = projection_matrix
+ if light_direction is not None:
+ if isinstance(light_direction, np.ndarray):
+ kwargs['lightDirection'] = light_direction.ravel().tolist()
+ else:
+ kwargs['lightDirection'] = light_direction
+ if light_color is not None:
+ if isinstance(light_color, np.ndarray):
+ kwargs['lightColor'] = light_color
+ else:
+ kwargs['lightColor'] = light_color
+ if light_distance is not None:
+ kwargs['lightDistance'] = light_distance
+ if shadow is not None:
+ kwargs['shadow'] = int(shadow)
+ if light_ambient_coeff is not None:
+ kwargs['lightAmbientCoeff'] = light_ambient_coeff
+ if light_diffuse_coeff is not None:
+ kwargs['lightDiffuseCoeff'] = light_diffuse_coeff
+ if light_specular_coeff is not None:
+ kwargs['lightSpecularCoeff'] = light_specular_coeff
+ if renderer is not None:
+ kwargs['renderer'] = renderer
+ if flags is not None:
+ kwargs['flags'] = flags
+
+ img = np.array(self.sim.getCameraImage(width, height, **kwargs)[4])
+ img = img.reshape(width, height)
+ return img
+
+ ##############
+ # Collisions #
+ ##############
+
+ def create_collision_shape(self, shape_type, radius=0.5, half_extents=(1., 1., 1.), height=1., filename=None,
+ mesh_scale=(1., 1., 1.), plane_normal=(0., 0., 1.), flags=-1,
+ collision_frame_position=None, collision_frame_orientation=None):
+ """
+ Create collision shape in the simulator.
+
+ Args:
+ shape_type (int): type of shape; GEOM_SPHERE (=2), GEOM_BOX (=3), GEOM_CAPSULE (=7), GEOM_CYLINDER (=4),
+ GEOM_PLANE (=6), GEOM_MESH (=5)
+ radius (float): only for GEOM_SPHERE, GEOM_CAPSULE, GEOM_CYLINDER
+ half_extents (np.float[3], list/tuple of 3 floats): only for GEOM_BOX.
+ height (float): only for GEOM_CAPSULE, GEOM_CYLINDER (length = height).
+ filename (str): Filename for GEOM_MESH, currently only Wavefront .obj. Will create convex hulls for each
+ object (marked as 'o') in the .obj file.
+ mesh_scale (np.float[3], list/tuple of 3 floats): scale of mesh (only for GEOM_MESH).
+ plane_normal (np.float[3], list/tuple of 3 floats): plane normal (only for GEOM_PLANE).
+ flags (int): unused / to be decided
+ collision_frame_position (np.float[3]): translational offset of the collision shape with respect to the
+ link frame
+ collision_frame_orientation (np.float[4]): rotational offset (quaternion x,y,z,w) of the collision shape
+ with respect to the link frame
+
+ Returns:
+ int: The return value is a non-negative int unique id for the collision shape or -1 if the call failed.
+ """
+ # add few variables
+ kwargs = {}
+ if collision_frame_position is not None:
+ kwargs['collisionFramePosition'] = collision_frame_position
+ if collision_frame_orientation is not None:
+ kwargs['collisionFrameOrientation'] = collision_frame_orientation
+
+ if shape_type == self.sim.GEOM_SPHERE:
+ return self.sim.createCollisionShape(shape_type, radius=radius, **kwargs)
+ elif shape_type == self.sim.GEOM_BOX:
+ return self.sim.createCollisionShape(shape_type, halfExtents=half_extents, **kwargs)
+ elif shape_type == self.sim.GEOM_CAPSULE or shape_type == self.sim.GEOM_CYLINDER:
+ return self.sim.createCollisionShape(shape_type, radius=radius, height=height, **kwargs)
+ elif shape_type == self.sim.GEOM_PLANE:
+ return self.sim.createCollisionShape(shape_type, planeNormal=plane_normal, **kwargs)
+ elif shape_type == self.sim.GEOM_MESH:
+ return self.sim.createCollisionShape(shape_type, fileName=filename, **kwargs)
+ else:
+ raise ValueError("Unknown collision shape type.")
+
+ def get_collision_shape_data(self, object_id, link_id=-1):
+ """
+ Get the collision shape data associated with the specified object id and link id.
+
+ Args:
+ object_id (int): object unique id.
+ link_id (int): link index or -1 for the base.
+
+ Returns:
+ int: object unique id.
+ int: link id.
+ int: geometry type; GEOM_BOX (=3), GEOM_SPHERE (=2), GEOM_CAPSULE (=7), GEOM_MESH (=5), GEOM_PLANE (=6)
+ np.float[3]: depends on geometry type:
+ for GEOM_BOX: extents,
+ for GEOM_SPHERE: dimensions[0] = radius,
+ for GEOM_CAPSULE and GEOM_CYLINDER: dimensions[0] = height (length), dimensions[1] = radius.
+ For GEOM_MESH: dimensions is the scaling factor.
+ str: Only for GEOM_MESH: file name (and path) of the collision mesh asset.
+ np.float[3]: Local position of the collision frame with respect to the center of mass/inertial frame
+ np.float[4]: Local orientation of the collision frame with respect to the inertial frame
+ """
+ object_id, link_id, geom_type, dimensions, filename, \
+ position, orientation = self.sim.getCollisionShapeData(object_id, link_id)
+ return object_id, link_id, geom_type, np.array(dimensions), filename, np.array(position), np.array(orientation)
+
+ def get_overlapping_objects(self, aabb_min, aabb_max):
+ """
+ This query will return all the unique ids of objects that have Axis Aligned Bounding Box (AABB) overlap with
+ a given axis aligned bounding box. Note that the query is conservative and may return additional objects that
+ don't have actual AABB overlap. This happens because the acceleration structures have some heuristic that
+ enlarges the AABBs a bit (extra margin and extruded along the velocity vector).
+
+ Args:
+ aabb_min (np.float[3]): minimum coordinates of the aabb
+ aabb_max (np.float[3]): maximum coordinates of the aabb
+
+ Returns:
+ list of int: list of object unique ids.
+ """
+ return self.sim.getOverlappingObjects(aabb_min, aabb_max)
+
+ def get_aabb(self, body_id, link_id=-1):
+ """
+ You can query the axis aligned bounding box (in world space) given an object unique id, and optionally a link
+ index. (when you don't pass the link index, or use -1, you get the AABB of the base).
+
+ Args:
+ body_id (int): object unique id as returned by creation methods
+ link_id (int): link index in range [0..`getNumJoints(..)]
+
+ Returns:
+ np.float[3]: minimum coordinates of the axis aligned bounding box
+ np.float[3]: maximum coordinates of the axis aligned bounding box
+ """
+ aabb_min, aabb_max = self.sim.getAABB(body_id, link_id)
+ return np.array(aabb_min), np.array(aabb_max)
+
+ def get_contact_points(self, body_a, body_b, link_id_a=None, link_id_b=None):
+ """
+ Returns the contact points computed during the most recent call to `step`.
+
+ Args:
+ body_a (int): only report contact points that involve body A
+ body_b (int): only report contact points that involve body B. Important: you need to have a valid body A
+ if you provide body B
+ link_id_a (int): only report contact points that involve link index of body A
+ link_id_b (int): only report contact points that involve link index of body B
+
+ Returns:
+ list:
+ int: contact flag (reserved)
+ int: body unique id of body A
+ int: body unique id of body B
+ int: link index of body A, -1 for base
+ int: link index of body B, -1 for base
+ np.float[3]: contact position on A, in Cartesian world coordinates
+ np.float[3]: contact position on B, in Cartesian world coordinates
+ np.float[3]: contact normal on B, pointing towards A
+ float: contact distance, positive for separation, negative for penetration
+ float: normal force applied during the last `step`
+ float: lateral friction force in the first lateral friction direction (see next returned value)
+ np.float[3]: first lateral friction direction
+ float: lateral friction force in the second lateral friction direction (see next returned value)
+ np.float[3]: second lateral friction direction
+ """
+ kwargs = {}
+ if body_a is not None:
+ kwargs['bodyA'] = body_a
+ if link_id_a is not None:
+ kwargs['linkIndexA'] = link_id_a
+ if body_b is not None:
+ kwargs['bodyB'] = body_b
+ if link_id_b is not None:
+ kwargs['linkIndexB'] = link_id_b
+
+ results = self.sim.getContactPoints(**kwargs)
+ if len(results) == 0:
+ return results
+ return [[r[0], r[1], r[2], r[3], r[4], np.array(r[5]), np.array(r[6]), np.array(r[7]), r[8], r[9], r[10],
+ np.array(r[11]), r[12], np.array(r[13])] for r in results]
+
+ def get_closest_points(self, body_a, body_b, distance, link_id_a=None, link_id_b=None):
+ """
+ Computes the closest points, independent from `step`. This also lets you compute closest points of objects
+ with an arbitrary separating distance. In this query there will be no normal forces reported.
+
+ Args:
+ body_a (int): only report contact points that involve body A
+ body_b (int): only report contact points that involve body B. Important: you need to have a valid body A
+ if you provide body B
+ distance (float): If the distance between objects exceeds this maximum distance, no points may be returned.
+ link_id_a (int): only report contact points that involve link index of body A
+ link_id_b (int): only report contact points that involve link index of body B
+
+ Returns:
+ list:
+ int: contact flag (reserved)
+ int: body unique id of body A
+ int: body unique id of body B
+ int: link index of body A, -1 for base
+ int: link index of body B, -1 for base
+ np.float[3]: contact position on A, in Cartesian world coordinates
+ np.float[3]: contact position on B, in Cartesian world coordinates
+ np.float[3]: contact normal on B, pointing towards A
+ float: contact distance, positive for separation, negative for penetration
+ float: normal force applied during the last `step`. Always equal to 0.
+ float: lateral friction force in the first lateral friction direction (see next returned value)
+ np.float[3]: first lateral friction direction
+ float: lateral friction force in the second lateral friction direction (see next returned value)
+ np.float[3]: second lateral friction direction
+ """
+ kwargs = {}
+ if link_id_a is not None:
+ kwargs['linkIndexA'] = link_id_a
+ if link_id_b is not None:
+ kwargs['linkIndexB'] = link_id_b
+
+ results = self.sim.getContactPoints(body_a, body_b, distance, **kwargs)
+ if len(results) == 0:
+ return results
+ return [[r[0], r[1], r[2], r[3], r[4], np.array(r[5]), np.array(r[6]), np.array(r[7]), r[8], r[9], r[10],
+ np.array(r[11]), r[12], np.array(r[13])] for r in results]
+
+ def ray_test(self, from_position, to_position):
+ """
+ Performs a single raycast to find the intersection information of the first object hit.
+
+ Args:
+ from_position (np.float[3]): start of the ray in world coordinates
+ to_position (np.float[3]): end of the ray in world coordinates
+
+ Returns:
+ int: object unique id of the hit object
+ int: link index of the hit object, or -1 if none/parent
+ float: hit fraction along the ray in range [0,1] along the ray.
+ np.float[3]: hit position in Cartesian world coordinates
+ np.float[3]: hit normal in Cartesian world coordinates
+ """
+ if isinstance(from_position, np.ndarray):
+ from_position = from_position.ravel().tolist()
+ if isinstance(to_position, np.ndarray):
+ to_position = to_position.ravel().tolist()
+ object_id, link_id, hit_fraction, position, normal = self.sim.rayTest(from_position, to_position)
+ return object_id, link_id, hit_fraction, np.array(position), np.array(normal)
+
+ def ray_test_batch(self, from_positions, to_positions, parent_object_id=None, parent_link_id=None):
+ """Perform a batch of raycasts to find the intersection information of the first objects hit.
+
+ This is similar to the rayTest, but allows you to provide an array of rays, for faster execution. The size of
+ 'rayFromPositions' needs to be equal to the size of 'rayToPositions'. You can one ray result per ray, even if
+ there is no intersection: you need to use the objectUniqueId field to check if the ray has hit anything: if
+ the objectUniqueId is -1, there is no hit. In that case, the 'hit fraction' is 1. The maximum number of rays
+ per batch is `pybullet.MAX_RAY_INTERSECTION_BATCH_SIZE`.
+
+ Args:
+ from_positions (np.array[N,3]): list of start points for each ray, in world coordinates
+ to_positions (np.array[N,3]): list of end points for each ray in world coordinates
+ parent_object_id (int): ray from/to is in local space of a parent object
+ parent_link_id (int): ray from/to is in local space of a parent object
+
+ Returns:
+ list:
+ int: object unique id of the hit object
+ int: link index of the hit object, or -1 if none/parent
+ float: hit fraction along the ray in range [0,1] along the ray.
+ np.float[3]: hit position in Cartesian world coordinates
+ np.float[3]: hit normal in Cartesian world coordinates
+ """
+ if isinstance(from_positions, np.ndarray):
+ from_positions = from_positions.tolist()
+ if isinstance(to_positions, np.ndarray):
+ to_positions = to_positions.tolist()
+
+ kwargs = {}
+ if parent_object_id is not None:
+ kwargs['parentObjectUniqueId'] = parent_object_id
+ if parent_link_id is not None:
+ kwargs['parentLinkIndex'] = parent_link_id
+
+ results = self.sim.rayTestBatch(from_positions, to_positions, **kwargs)
+ if len(results) == 0:
+ return results
+ return [[r[0], r[1], r[2], np.array(r[3]), np.array(r[4])] for r in results]
+
+ def set_collision_filter_group_mask(self, body_id, link_id, filter_group, filter_mask):
+ """
+ Enable/disable collision detection between groups of objects. Each body is part of a group. It collides with
+ other bodies if their group matches the mask, and vise versa. The following check is performed using the group
+ and mask of the two bodies involved. It depends on the collision filter mode.
+
+ Args:
+ body_id (int): unique id of the body to be configured
+ link_id (int): link index of the body to be configured
+ filter_group (int): bitwise group of the filter
+ filter_mask (int): bitwise mask of the filter
+ """
+ self.sim.setCollisionFilterGroupMask(body_id, link_id, filter_group, filter_mask)
+
+ def set_collision_filter_pair(self, body_a, body_b, link_a=-1, link_b=-1, enable=True):
+ """
+ Enable/disable collision between two bodies/links.
+
+ Args:
+ body_a (int): unique id of body A to be filtered
+ body_b (int): unique id of body B to be filtered, A==B implies self-collision
+ link_a (int): link index of body A
+ link_b (int): link index of body B
+ enable (bool): True to enable collision, False to disable collision
+ """
+ self.sim.setCollisionFilterPair(body_a, body_b, link_a, link_b, int(enable))
+
+ ###########################
+ # Kinematics and Dynamics #
+ ###########################
+
+ def get_dynamics_info(self, body_id, link_id=-1):
+ """
+ Get dynamic information about the mass, center of mass, friction and other properties of the base and links.
+
+ Args:
+ body_id (int): body/object unique id.
+ link_id (int): link/joint index or -1 for the base.
+
+ Returns:
+ float: mass in kg
+ float: friction coefficient
+ np.float[3]: local inertia diagonal. Note that links and base are centered around the center of mass and
+ aligned with the principal axes of inertia.
+ np.float[3]: position of inertial frame in local coordinates of the joint frame
+ np.float[4]: orientation of inertial frame in local coordinates of joint frame
+ float: coefficient of restitution
+ float: rolling friction coefficient orthogonal to contact normal
+ float: spinning friction coefficient around contact normal
+ float: damping of contact constraints. -1 if not available.
+ float: stiffness of contact constraints. -1 if not available.
+ """
+ mass, friction, inertia, pos, quat, restitution, roll, spin, kd, kp = self.sim.getDynamicsInfo(body_id, link_id)
+ return mass, friction, np.array(inertia), np.array(pos), np.array(quat), restitution, roll, spin, kd, kp
+
+ def change_dynamics(self, body_id, link_id=-1, mass=None, lateral_friction=None, spinning_friction=None,
+ rolling_friction=None, restitution=None, linear_damping=None, angular_damping=None,
+ contact_stiffness=None, contact_damping=None, friction_anchor=None,
+ local_inertia_diagonal=None, joint_damping=None):
+ """
+ Change dynamic properties such as mass, friction and restitution coefficients .
+
+ Args:
+ body_id (int): object unique id, as returned by `load_urdf`, etc.
+ link_id (int): link index or -1 for the base.
+ mass (float): change the mass of the link (or base for link index -1)
+ lateral_friction (float): lateral (linear) contact friction
+ spinning_friction (float): torsional friction around the contact normal
+ rolling_friction (float): torsional friction orthogonal to contact normal
+ restitution (float): bouncyness of contact. Keep it a bit less than 1.
+ linear_damping (float): linear damping of the link (0.04 by default)
+ angular_damping (float): angular damping of the link (0.04 by default)
+ contact_stiffness (float): stiffness of the contact constraints, used together with `contact_damping`
+ contact_damping (float): damping of the contact constraints for this body/link. Used together with
+ `contact_stiffness`. This overrides the value if it was specified in the URDF file in the contact
+ section.
+ friction_anchor (int): enable or disable a friction anchor: positional friction correction (disabled by
+ default, unless set in the URDF contact section)
+ local_inertia_diagonal (np.float[3]): diagonal elements of the inertia tensor. Note that the base and
+ links are centered around the center of mass and aligned with the principal axes of inertia so there
+ are no off-diagonal elements in the inertia tensor.
+ joint_damping (float): joint damping coefficient applied at each joint. This coefficient is read from URDF
+ joint damping field. Keep the value close to 0.
+ `joint_damping_force = -damping_coefficient * joint_velocity`.
+ """
+ kwargs = {}
+ if mass is not None:
+ kwargs['mass'] = mass
+ if lateral_friction is not None:
+ kwargs['lateralFriction'] = lateral_friction
+ if spinning_friction is not None:
+ kwargs['spinningFriction'] = spinning_friction
+ if rolling_friction is not None:
+ kwargs['rollingFriction'] = rolling_friction
+ if restitution is not None:
+ kwargs['restitution'] = restitution
+ if linear_damping is not None:
+ kwargs['linearDamping'] = linear_damping
+ if angular_damping is not None:
+ kwargs['angularDamping'] = angular_damping
+ if contact_stiffness is not None:
+ kwargs['contactStiffness'] = contact_stiffness
+ if contact_damping is not None:
+ kwargs['contactDamping'] = contact_damping
+ if friction_anchor is not None:
+ kwargs['frictionAnchor'] = friction_anchor
+ if local_inertia_diagonal is not None:
+ kwargs['localInertiaDiagonal'] = local_inertia_diagonal
+ if joint_damping is not None:
+ kwargs['jointDamping'] = joint_damping
+
+ self.sim.changeDynamics(body_id, link_id, **kwargs)
+
+ def calculate_jacobian(self, body_id, link_id, local_position, q, dq, des_ddq):
+ """
+ Return the full geometric Jacobian matrix :math:`J(q) = [J_{lin}(q), J_{ang}(q)]^T`, such that:
+
+ .. math:: v = [\dot{p}, \omega]^T = J(q) \dot{q}
+
+ where :math:`\dot{p}` is the Cartesian linear velocity of the link, and :math:`\omega` is its angular velocity.
+
+ Warnings: if we have a floating base then the Jacobian will also include columns corresponding to the root
+ link DoFs (at the beginning). If it is a fixed base, it will only have columns associated with the joints.
+
+ Args:
+ body_id (int): unique body id.
+ link_id (int): link id.
+ local_position (np.float[3]): the point on the specified link to compute the Jacobian (in link local
+ coordinates around its center of mass). If None, it will use the CoM position (in the link frame).
+ q (np.float[N]): joint positions of size N, where N is the number of DoFs.
+ dq (np.float[N]): joint velocities of size N, where N is the number of DoFs.
+ des_ddq (np.float[N]): desired joint accelerations of size N.
+
+ Returns:
+ np.float[6,N], np.float[6,(6+N)]: full geometric (linear and angular) Jacobian matrix. The number of
+ columns depends if the base is fixed or floating.
+ """
+ # Note that q, dq, ddq have to be lists in PyBullet (it doesn't work with numpy arrays)
+ if isinstance(local_position, np.ndarray):
+ local_position = local_position.ravel().tolist()
+ if isinstance(q, np.ndarray):
+ q = q.ravel().tolist()
+ if isinstance(dq, np.ndarray):
+ dq = dq.ravel().tolist()
+ if isinstance(des_ddq, np.ndarray):
+ des_ddq = des_ddq.ravel().tolist()
+
+ # calculate full jacobian
+ lin_jac, ang_jac = self.sim.calculateJacobian(body_id, link_id, localPosition=local_position,
+ objPositions=q, objVelocities=dq, objAccelerations=des_ddq)
+
+ return np.vstack((lin_jac, ang_jac))
+
+ def calculate_mass_matrix(self, body_id, q):
+ """
+ Return the mass/inertia matrix :math:`H(q)`, which is used in the rigid-body equation of motion (EoM) in joint
+ space given by (see [1]):
+
+ .. math:: \tau = H(q)\ddot{q} + C(q,\dot{q})
+
+ where :math:`\tau` is the vector of applied torques, :math:`H(q)` is the inertia matrix, and
+ :math:`C(q,\dot{q}) \dot{q}` is the vector accounting for Coriolis, centrifugal forces, gravity, and any
+ other forces acting on the system except the applied torques :math:`\tau`.
+
+ Warnings: If the base is floating, it will return a [6+N,6+N] inertia matrix, where N is the number of actuated
+ joints. If the base is fixed, it will return a [N,N] inertia matrix
+
+ Args:
+ body_id (int): body unique id.
+ q (np.float[N]): joint positions of size N, where N is the total number of DoFs.
+
+ Returns:
+ np.float[N,N], np.float[6+N,6+N]: inertia matrix
+ """
+ if isinstance(q, np.ndarray):
+ q = q.ravel().tolist()
+ return np.array(self.sim.calculateMassMatrix(body_id, q))
+
+ def calculate_inverse_kinematics(self, body_id, link_id, position, orientation=None, lower_limits=None,
+ upper_limits=None, joint_ranges=None, rest_poses=None, joint_dampings=None,
+ solver=None, q_curr=None, max_iters=None, threshold=None):
+ """
+ Compute the FULL Inverse kinematics; it will return a position for all the actuated joints.
+
+ "You can compute the joint angles that makes the end-effector reach a given target position in Cartesian world
+ space. Internally, Bullet uses an improved version of Samuel Buss Inverse Kinematics library. At the moment
+ only the Damped Least Squares method with or without Null Space control is exposed, with a single end-effector
+ target. Optionally you can also specify the target orientation of the end effector. In addition, there is an
+ option to use the null-space to specify joint limits and rest poses. This optional null-space support requires
+ all 4 lists (lower_limits, upper_limits, joint_ranges, rest_poses), otherwise regular IK will be used." [1]
+
+ Args:
+ body_id (int): body unique id, as returned by `load_urdf`, etc.
+ link_id (int): end effector link index.
+ position (np.float[3]): target position of the end effector (its link coordinate, not center of mass
+ coordinate!). By default this is in Cartesian world space, unless you provide `q_curr` joint angles.
+ orientation (np.float[4]): target orientation in Cartesian world space, quaternion [x,y,w,z]. If not
+ specified, pure position IK will be used.
+ lower_limits (np.float[N], list of N floats): lower joint limits. Optional null-space IK.
+ upper_limits (np.float[N], list of N floats): upper joint limits. Optional null-space IK.
+ joint_ranges (np.float[N], list of N floats): range of value of each joint.
+ rest_poses (np.float[N], list of N floats): joint rest poses. Favor an IK solution closer to a given rest
+ pose.
+ joint_dampings (np.float[N], list of N floats): joint damping factors. Allow to tune the IK solution using
+ joint damping factors.
+ solver (int): p.IK_DLS (=0) or p.IK_SDLS (=1), Damped Least Squares or Selective Damped Least Squares, as
+ described in the paper by Samuel Buss "Selectively Damped Least Squares for Inverse Kinematics".
+ q_curr (np.float[N]): list of joint positions. By default PyBullet uses the joint positions of the body.
+ If provided, the targetPosition and targetOrientation is in local space!
+ max_iters (int): maximum number of iterations. Refine the IK solution until the distance between target
+ and actual end effector position is below this threshold, or the `max_iters` is reached.
+ threshold (float): residual threshold. Refine the IK solution until the distance between target and actual
+ end effector position is below this threshold, or the `max_iters` is reached.
+
+ Returns:
+ np.float[N]: joint positions (for each actuated joint).
+ """
+ kwargs = {}
+ if orientation is not None:
+ if isinstance(orientation, np.ndarray):
+ orientation = orientation.ravel().tolist()
+ kwargs['targetOrientation'] = orientation
+ if lower_limits is not None and upper_limits is not None and joint_ranges is not None and \
+ rest_poses is not None:
+ kwargs['lowerLimits'], kwargs['upperLimits'] = lower_limits, upper_limits
+ kwargs['jointRanges'], kwargs['restPoses'] = joint_ranges, rest_poses
+
+ if q_curr is not None:
+ if isinstance(q_curr, np.ndarray):
+ q_curr = q_curr.ravel().tolist()
+ kwargs['currentPosition'] = q_curr
+ if joint_dampings is not None:
+ if isinstance(joint_dampings, np.ndarray):
+ joint_dampings = joint_dampings.ravel().tolist()
+ kwargs['jointDamping'] = joint_dampings
+
+ if solver is not None:
+ kwargs['solver'] = solver
+ if max_iters is not None:
+ kwargs['maxNumIterations'] = max_iters
+ if threshold is not None:
+ kwargs['residualThreshold'] = threshold
+
+ return np.array(self.sim.calculateInverseKinematics(body_id, link_id, position, **kwargs))
+
+ def calculate_inverse_dynamics(self, body_id, q, dq, des_ddq):
+ r"""
+ Starting from the specified joint positions :math:`q` and velocities :math:`\dot{q}`, it computes the joint
+ torques :math:`\tau` required to reach the desired joint accelerations :math:`\ddot{q}_{des}`. That is,
+ :math:`\tau = ID(model, q, \dot{q}, \ddot{q}_{des})`.
+
+ Specifically, it uses the rigid-body equation of motion in joint space given by (see [1]):
+
+ .. math:: \tau = H(q)\ddot{q} + C(q,\dot{q})
+
+ where :math:`\tau` is the vector of applied torques, :math:`H(q)` is the inertia matrix, and
+ :math:`C(q,\dot{q}) \dot{q}` is the vector accounting for Coriolis, centrifugal forces, gravity, and any
+ other forces acting on the system except the applied torques :math:`\tau`.
+
+ Normally, a more popular form of this equation of motion (in joint space) is given by:
+
+ .. math:: H(q) \ddot{q} + S(q,\dot{q}) \dot{q} + g(q) = \tau + J^T(q) F
+
+ which is the same as the first one with :math:`C = S\dot{q} + g(q) - J^T(q) F`. However, this last formulation
+ is useful to understand what happens when we set some variables to 0.
+ Assuming that there are no forces acting on the system, and giving desired joint accelerations of 0, this
+ method will return :math:`\tau = S(q,\dot{q}) \dot{q} + g(q)`. If in addition joint velocities are also 0,
+ it will return :math:`\tau = g(q)` which can for instance be useful for gravity compensation.
+
+ For forward dynamics, which computes the joint accelerations given the joint positions, velocities, and
+ torques (that is, :math:`\ddot{q} = FD(model, q, \dot{q}, \tau)`, this can be computed using
+ :math:`\ddot{q} = H^{-1} (\tau - C)` (see also `computeFullFD`). For more information about different
+ control schemes (position, force, impedance control and others), or about the formulation of the equation
+ of motion in task/operational space (instead of joint space), check the references [1-4].
+
+ Args:
+ body_id (int): body unique id.
+ q (np.float[N]): joint positions
+ dq (np.float[N]): joint velocities
+ des_ddq (np.float[N]): desired joint accelerations
+
+ Returns:
+ np.float[N]: joint torques computed using the rigid-body equation of motion
+
+ References:
+ [1] "Rigid Body Dynamics Algorithms", Featherstone, 2008, chap1.1
+ [2] "Robotics: Modelling, Planning and Control", Siciliano et al., 2010
+ [3] "Springer Handbook of Robotics", Siciliano et al., 2008
+ [4] Lecture on "Impedance Control" by Prof. De Luca, Universita di Roma,
+ http://www.diag.uniroma1.it/~deluca/rob2_en/15_ImpedanceControl.pdf
+ """
+ # convert numpy arrays to lists
+ if isinstance(q, np.ndarray):
+ q = q.ravel().tolist()
+ if isinstance(dq, np.ndarray):
+ dq = dq.ravel().tolist()
+ if isinstance(des_ddq, np.ndarray):
+ des_ddq = des_ddq.ravel().tolist()
+
+ # return the joint torques to be applied for the desired joint accelerations
+ return np.array(self.sim.calculateInverseDynamics(body_id, q, dq, des_ddq))
+
+ def calculate_forward_dynamics(self, body_id, q, dq, torques):
+ r"""
+ Given the specified joint positions :math:`q` and velocities :math:`\dot{q}`, and joint torques :math:`\tau`,
+ it computes the joint accelerations :math:`\ddot{q}`. That is, :math:`\ddot{q} = FD(model, q, \dot{q}, \tau)`.
+
+ Specifically, it uses the rigid-body equation of motion in joint space given by (see [1]):
+
+ .. math:: \ddot{q} = H(q)^{-1} (\tau - C(q,\dot{q}))
+
+ where :math:`\tau` is the vector of applied torques, :math:`H(q)` is the inertia matrix, and
+ :math:`C(q,\dot{q}) \dot{q}` is the vector accounting for Coriolis, centrifugal forces, gravity, and any
+ other forces acting on the system except the applied torques :math:`\tau`.
+
+ Normally, a more popular form of this equation of motion (in joint space) is given by:
+
+ .. math:: H(q) \ddot{q} + S(q,\dot{q}) \dot{q} + g(q) = \tau + J^T(q) F
+
+ which is the same as the first one with :math:`C = S\dot{q} + g(q) - J^T(q) F`. However, this last formulation
+ is useful to understand what happens when we set some variables to 0.
+ Assuming that there are no forces acting on the system, and giving desired joint torques of 0, this
+ method will return :math:`\ddot{q} = - H(q)^{-1} (S(q,\dot{q}) \dot{q} + g(q))`. If in addition
+ the joint velocities are also 0, it will return :math:`\ddot{q} = - H(q)^{-1} g(q)` which are
+ the accelerations due to gravity.
+
+ For inverse dynamics, which computes the joint torques given the joint positions, velocities, and
+ accelerations (that is, :math:`\tau = ID(model, q, \dot{q}, \ddot{q})`, this can be computed using
+ :math:`\tau = H(q)\ddot{q} + C(q,\dot{q})`. For more information about different
+ control schemes (position, force, impedance control and others), or about the formulation of the equation
+ of motion in task/operational space (instead of joint space), check the references [1-4].
+
+ Args:
+ body_id (int): unique body id.
+ q (np.float[N]): joint positions
+ dq (np.float[N]): joint velocities
+ torques (np.float[N]): desired joint torques
+
+ Returns:
+ float[N]: joint accelerations computed using the rigid-body equation of motion
+
+ References:
+ [1] "Rigid Body Dynamics Algorithms", Featherstone, 2008, chap1.1
+ [2] "Robotics: Modelling, Planning and Control", Siciliano et al., 2010
+ [3] "Springer Handbook of Robotics", Siciliano et al., 2008
+ [4] Lecture on "Impedance Control" by Prof. De Luca, Universita di Roma,
+ http://www.diag.uniroma1.it/~deluca/rob2_en/15_ImpedanceControl.pdf
+ """
+ # convert numpy arrays to lists
+ if isinstance(q, np.ndarray):
+ q = q.ravel().tolist()
+
+ # compute and return joint accelerations
+ torques = np.array(torques)
+ Hinv = np.linalg.inv(self.calculate_mass_matrix(body_id, q))
+ C = self.calculate_inverse_dynamics(body_id, q, dq, np.zeros(len(q)))
+ acc = Hinv.dot(torques - C)
+ return acc
+
+ #########
+ # Debug #
+ #########
+
+ def add_user_debug_line(self, from_pos, to_pos, rgb_color=None, width=None, lifetime=None, parent_object_id=None,
+ parent_link_id=None, line_id=None):
+ """Add a user debug line in the simulator.
+
+ You can add a 3d line specified by a 3d starting point (from) and end point (to), a color [red,green,blue],
+ a line width and a duration in seconds.
+
+ Args:
+ from_pos (np.float[3]): starting point of the line in Cartesian world coordinates
+ to_pos (np.float[3]): end point of the line in Cartesian world coordinates
+ rgb_color (np.float[3]): RGB color (each channel in range [0,1])
+ width (float): line width (limited by OpenGL implementation).
+ lifetime (float): use 0 for permanent line, or positive time in seconds (afterwards the line with be
+ removed automatically)
+ parent_object_id (int): draw line in local coordinates of a parent object.
+ parent_link_id (int): draw line in local coordinates of a parent link.
+ line_id (int): replace an existing line item (to avoid flickering of remove/add).
+
+ Returns:
+ int: unique user debug line id.
+ """
+ kwargs = {}
+ if rgb_color is not None:
+ kwargs['lineColorRGB'] = rgb_color
+ if width is not None:
+ kwargs['lineWidth'] = width
+ if lifetime is not None:
+ kwargs['lifeTime'] = lifetime
+ if parent_object_id is not None:
+ kwargs['parentObjectUniqueId'] = parent_object_id
+ if parent_link_id is not None:
+ kwargs['parentLinkIndex'] = parent_link_id
+ if line_id is not None:
+ kwargs['replaceItemUniqueId'] = line_id
+
+ return self.sim.addUserDebugLine(lineFromXYZ=from_pos, lineToXYZ=to_pos, **kwargs)
+
+ def add_user_debug_text(self, text, position, rgb_color=None, size=None, lifetime=None, orientation=None,
+ parent_object_id=None, parent_link_id=None, text_id=None):
+ """
+ Add 3D text at a specific location using a color and size.
+
+ Args:
+ text (str): text.
+ position (np.float[3]): 3d position of the text in Cartesian world coordinates.
+ rgb_color (list/tuple of 3 floats): RGB color; each component in range [0..1]
+ size (float): text size
+ lifetime (float): use 0 for permanent text, or positive time in seconds (afterwards the text with be
+ removed automatically)
+ orientation (np.float[4]): By default, debug text will always face the camera, automatically rotation.
+ By specifying a text orientation (quaternion), the orientation will be fixed in world space or local
+ space (when parent is specified). Note that a different implementation/shader is used for camera
+ facing text, with different appearance: camera facing text uses bitmap fonts, text with specified
+ orientation uses TrueType font.
+ parent_object_id (int): draw text in local coordinates of a parent object.
+ parent_link_id (int): draw text in local coordinates of a parent link.
+ text_id (int): replace an existing text item (to avoid flickering of remove/add).
+
+ Returns:
+ int: unique user debug text id.
+ """
+ kwargs = {}
+ if rgb_color is not None:
+ kwargs['textColorRGB'] = rgb_color
+ if size is not None:
+ kwargs['textSize'] = size
+ if lifetime is not None:
+ kwargs['lifeTime'] = lifetime
+ if orientation is not None:
+ kwargs['textOrientation'] = orientation
+ if parent_object_id is not None:
+ kwargs['parentObjectUniqueId'] = parent_object_id
+ if parent_link_id is not None:
+ kwargs['parentLinkIndex'] = parent_link_id
+ if text_id is not None:
+ kwargs['replaceItemUniqueId'] = text_id
+
+ return self.sim.addUserDebugText(text=text, textPosition=position, **kwargs)
+
+ def add_user_debug_parameter(self, name, min_range, max_range, start_value):
+ """
+ Add custom sliders to tune parameters.
+
+ Args:
+ name (str): name of the parameter.
+ min_range (float): minimum value.
+ max_range (float): maximum value.
+ start_value (float): starting value.
+
+ Returns:
+ int: unique user debug parameter id.
+ """
+ return self.sim.addUserDebugParameter(paramName=name, rangeMin=min_range, rangeMax=max_range,
+ startValue=start_value)
+
+ def read_user_debug_parameter(self, parameter_id):
+ """
+ Read the value of the parameter / slider.
+
+ Args:
+ parameter_id: unique user debug parameter id.
+
+ Returns:
+ float: reading of the parameter.
+ """
+ return self.sim.readUserDebugParameter(parameter_id)
+
+ def remove_user_debug_item(self, item_id):
+ """
+ Remove the specified user debug item (line, text, parameter) from the simulator.
+
+ Args:
+ item_id (int): unique id of the debug item to be removed (line, text etc)
+ """
+ self.sim.removeUserDebugItem(item_id)
+
+ def remove_all_user_debug_items(self):
+ """
+ Remove all user debug items from the simulator.
+ """
+ self.sim.removeAllUserDebugItems()
+
+ def set_debug_object_color(self, object_id, link_id, rgb_color=(1, 0, 0)):
+ """
+ Override the color of a specific object and link.
+
+ Args:
+ object_id (int): unique object id.
+ link_id (int): link id.
+ rgb_color (float[3]): RGB debug color.
+ """
+ self.sim.setDebugObjectColor(object_id, link_id, rgb_color)
+
+ def add_user_data(self, object_id, key, value):
+ """
+ Add user data (at the moment text strings) attached to any link of a body. You can also override a previous
+ given value. You can add multiple user data to the same body/link.
+
+ Args:
+ object_id (int): unique object/link id.
+ key (str): key string.
+ value (str): value string.
+
+ Returns:
+ int: user data id.
+ """
+ return self.sim.addUserData(object_id, key, value)
+
+ def num_user_data(self, object_id):
+ """
+ Return the number of user data associated with the specified object/link id.
+
+ Args:
+ object_id (int): unique object/link id.
+
+ Returns:
+ int: the number of user data
+ """
+ return self.sim.getNumUserData(object_id)
+
+ def get_user_data(self, user_data_id):
+ """
+ Get the specified user data value.
+
+ Args:
+ user_data_id (int): unique user data id.
+
+ Returns:
+ str: value string.
+ """
+ return self.sim.getUserData(user_data_id)
+
+ def get_user_data_id(self, object_id, key):
+ """
+ Get the specified user data id.
+
+ Args:
+ object_id (int): unique object/link id.
+ key (str): key string.
+
+ Returns:
+ int: user data id.
+ """
+ return self.sim.getUserDataId(object_id, key)
+
+ def get_user_data_info(self, object_id, index):
+ """
+ Get the user data info associated with the given object and index.
+
+ Args:
+ object_id (int): unique object id.
+ index (int): index (should be between [0, self.num_user_data(object_id)]).
+
+ Returns:
+ int: user data id.
+ str: key.
+ int: body id.
+ int: link index
+ int: visual shape index.
+ """
+ return self.sim.getUserDataInfo(object_id, index)
+
+ def remove_user_data(self, user_data_id):
+ """
+ Remove the specified user data.
+
+ Args:
+ user_data_id (int): user data id.
+ """
+ self.sim.removeUserData(user_data_id)
+
+ def sync_user_data(self):
+ """
+ Synchronize the user data.
+ """
+ self.sim.syncUserData()
+
+ def configure_debug_visualizer(self, flag, enable):
+ """Configure the debug visualizer camera.
+
+ Configure some settings of the built-in OpenGL visualizer, such as enabling or disabling wireframe,
+ shadows and GUI rendering.
+
+ Args:
+ flag (int): The feature to enable or disable, such as
+ COV_ENABLE_WIREFRAME (=3): show/hide the collision wireframe
+ COV_ENABLE_SHADOWS (=2): show/hide shadows
+ COV_ENABLE_GUI (=1): enable/disable the GUI
+ COV_ENABLE_VR_PICKING (=5): enable/disable VR picking
+ COV_ENABLE_VR_TELEPORTING (=4): enable/disable VR teleporting
+ COV_ENABLE_RENDERING (=7): enable/disable rendering
+ COV_ENABLE_TINY_RENDERER (=12): enable/disable tiny renderer
+ COV_ENABLE_VR_RENDER_CONTROLLERS (=6): render VR controllers
+ COV_ENABLE_KEYBOARD_SHORTCUTS (=9): enable/disable keyboard shortcuts
+ COV_ENABLE_MOUSE_PICKING (=10): enable/disable mouse picking
+ COV_ENABLE_Y_AXIS_UP (Z is default world up axis) (=11): enable/disable Y axis up
+ COV_ENABLE_RGB_BUFFER_PREVIEW (=13): enable/disable RGB buffer preview
+ COV_ENABLE_DEPTH_BUFFER_PREVIEW (=14): enable/disable Depth buffer preview
+ COV_ENABLE_SEGMENTATION_MARK_PREVIEW (=15): enable/disable segmentation mark preview
+ enable (bool): False (disable) or True (enable)
+ """
+ self.sim.configureDebugVisualizer(flag, int(enable))
+
+ def get_debug_visualizer(self):
+ """Get information about the debug visualizer camera.
+
+ Returns:
+ float: width of the visualizer camera
+ float: height of the visualizer camera
+ np.float[4,4]: view matrix [4,4]
+ np.float[4,4]: perspective projection matrix [4,4]
+ np.float[3]: camera up vector expressed in the Cartesian world space
+ np.float[3]: forward axis of the camera expressed in the Cartesian world space
+ np.float[3]: This is a horizontal vector that can be used to generate rays (for mouse picking or creating
+ a simple ray tracer for example)
+ np.float[3]: This is a vertical vector that can be used to generate rays (for mouse picking or creating a
+ simple ray tracer for example)
+ float: yaw angle (in radians) of the camera, in Cartesian local space coordinates
+ float: pitch angle (in radians) of the camera, in Cartesian local space coordinates
+ float: distance between the camera and the camera target
+ np.float[3]: target of the camera, in Cartesian world space coordinates
+ """
+ width, height, view, proj, up_vec, forward_vec,\
+ horizontal, vertical, yaw, pitch, dist, target = self.sim.getDebugVisualizerCamera()
+
+ # convert data to the correct data type
+ view = np.array(view).reshape(4, 4).T
+ proj = np.array(proj).reshape(4, 4).T
+ up_vec = np.array(up_vec)
+ forward_vec = np.array(forward_vec)
+ horizontal = np.array(horizontal)
+ vertical = np.array(vertical)
+ target = np.array(target)
+ yaw = np.deg2rad(yaw)
+ pitch = np.deg2rad(pitch)
+
+ # return the data
+ return width, height, view, proj, up_vec, forward_vec, horizontal, vertical, yaw, pitch, dist, target
+
+ def reset_debug_visualizer(self, distance, yaw, pitch, target_position):
+ """Reset the debug visualizer camera.
+
+ Reset the 3D OpenGL debug visualizer camera distance (between eye and camera target position), camera yaw and
+ pitch and camera target position
+
+ Args:
+ distance (float): distance from eye to camera target position
+ yaw (float): camera yaw angle (in radians) left/right
+ pitch (float): camera pitch angle (in radians) up/down
+ target_position (np.float[3]): target focus point of the camera
+ """
+ self.sim.resetDebugVisualizerCamera(cameraDistance=distance, cameraYaw=np.rad2deg(yaw),
+ cameraPitch=np.rad2deg(pitch), cameraTargetPosition=target_position)
+
+ ############################
+ # Events (mouse, keyboard) #
+ ############################
+
+ def get_keyboard_events(self):
+ """Get the key events.
+
+ Returns:
+ dict: {keyId: keyState}
+ * `keyID` is an integer (ascii code) representing the key. Some special keys like shift, arrows,
+ and others are are defined in pybullet such as `B3G_SHIFT`, `B3G_LEFT_ARROW`, `B3G_UP_ARROW`,...
+ * `keyState` is an integer. 3 if the button has been pressed, 1 if the key is down, 2 if the key has
+ been triggered.
+ """
+ return self.sim.getKeyboardEvents()
+
+ def get_mouse_events(self):
+ """Get the mouse events.
+
+ Returns:
+ list of mouse events:
+ eventType (int): 1 if the mouse is moving, 2 if a button has been pressed or released
+ mousePosX (float): x-coordinates of the mouse pointer
+ mousePosY (float): y-coordinates of the mouse pointer
+ buttonIdx (int): button index for left/middle/right mouse button. It is -1 if nothing,
+ 0 if left button, 1 if scroll wheel (pressed), 2 if right button
+ buttonState (int): 0 if nothing, 3 if the button has been pressed, 4 is the button has been released,
+ 1 if the key is down (never observed), 2 if the key has been triggered (never
+ observed).
+ """
+ return self.sim.getMouseEvents()
+
+ def get_mouse_and_keyboard_events(self):
+ """Get the mouse and key events.
+
+ Returns:
+ list: list of mouse events
+ dict: dictionary of key events
+ """
+ return self.sim.getMouseEvents(), self.sim.getKeyboardEvents()
+
+
+# Tests
+if __name__ == "__main__":
+ pass
+ # import inspect
+ # import pybullet as p
+ #
+ # f = p.createCollisionShape
+ # print(f)
+ # print(dir(f))
+ # specs = inspect.getargspec(f)
+ # print(zip(specs.args[-len(specs.defaults):], specs.defaults))
+ # print('')
+ #
+ # exit()
+ # sim = Bullet()
+ # s = sim.sim._client
+ # print('')
+ # # print(sim.get_physics_properties())
+ # f = s.createCollisionShape
+ # print(f)
+ # print(f.keywords)
+ # print(dir(f))
+ # specs = inspect.getargspec(f)
+ # print(zip(specs.args[-len(specs.defaults):], specs.defaults))
+ # print('')
diff --git a/pyrobolearn/simulators/bullet_ros.py b/pyrobolearn/simulators/bullet_ros.py
new file mode 100644
index 0000000..15ca7fc
--- /dev/null
+++ b/pyrobolearn/simulators/bullet_ros.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+"""Define the Bullet-ROS Simulator API.
+
+This is the main interface that communicates with the PyBullet simulator [1] and use ROS [3] to query the state of the
+robot and send instructions to it. By defining this interface, it allows to decouple the PyRoboLearn framework from
+the simulator. It also converts some data types to the ones required by PyBullet. For instance, some methods in
+PyBullet do not accepts numpy arrays but only lists. The interface provided here makes the necessary conversions.
+Using ROS to query the state of the robot, it changes the state of the robot in the simulator, and moving the robot
+in the simulator results in the real robot to move. Virtual sensors and actuators can also be defined.
+
+The signature of each method defined here are inspired by [1] but in accordance with the PEP8 style guide [2].
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+* `pyrobolearn.simulators.bullet.Bullet`
+* `pyrobolearn.simulators.ros.ROS`
+
+References:
+ [1] PyBullet: https://pybullet.org
+ [2] PyBullet Quickstart Guide: https://docs.google.com/document/d/10sXEhzFRSnvFcl3XxNGhnD4N2SedqwdAvK3dsihxVUA
+ [3] ROS: http://www.ros.org/
+ [4] PEP8: https://www.python.org/dev/peps/pep-0008/
+"""
+
+from simulator import Simulator
+from bullet import Bullet
+from ros import ROS
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class BulletROS(Simulator): # Bullet, ROS):
+ r"""Bullet ROS
+
+ Update the Bullet simulator based on the real robot(s): it updates the robot kinematic and dynamic state based on
+ the values returned from the real robot(s).
+
+ This can be useful for debug (check the differences between the real world and the simulated world), for virtual
+ sensors, actuators, and forces, to map the real world to the simulated one, etc.
+ """
+
+ def __init__(self):
+ super(BulletROS, self).__init__()
+ raise NotImplementedError
diff --git a/pyrobolearn/simulators/flex.py b/pyrobolearn/simulators/flex.py
new file mode 100644
index 0000000..1f91cdb
--- /dev/null
+++ b/pyrobolearn/simulators/flex.py
@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+"""Define the Nvidia FleX Simulator API.
+
+This is the main interface that communicates with the FleX simulator [1]. By defining this interface, it allows to
+decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+FleX.
+
+Warnings: We are waiting for [3] to publish their code.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] Nvidia FleX: https://developer.nvidia.com/flex
+ [2] Python bindings for the Nvidia FleX simulator: https://github.com/henryclever/FleX_PyBind11
+ [3] "GPU-Accelerated Robotic Simulation for Distributed Reinforcement Learning":
+ https://sites.google.com/view/accelerated-gpu-simulation/home
+"""
+
+from simulator import Simulator
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class Flex(Simulator):
+ r"""FleX simulator
+ """
+
+ def __init__(self):
+ super(Flex, self).__init__()
+ raise NotImplementedError
diff --git a/pyrobolearn/simulators/gazebo-ros.py b/pyrobolearn/simulators/gazebo-ros.py
new file mode 100644
index 0000000..a49e998
--- /dev/null
+++ b/pyrobolearn/simulators/gazebo-ros.py
@@ -0,0 +1,510 @@
+#!/usr/bin/env python
+"""Gazebo ROS simulator
+
+This simulator uses Gazebo as the simulator, ROS to communicate with this simulator (to send and receive any
+information related to the simulator and the objects inside of it like the robots), and RBDL to compute the kinematics
+and dynamics of the robots.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.ros_rbdl.ROS_RBDL`
+
+References:
+ [1] ROS: http://www.ros.org/
+ [2] Gazebo: http://gazebosim.org/
+ [3] RBDL: https://rbdl.bitbucket.io/
+"""
+
+import numpy as np
+import subprocess, os, signal, sys, time
+
+# import ROS and RBDL
+import rospy
+import rbdl
+
+# messages and services
+import std_msgs.msg as stdmsg
+import std_srvs.srv as stdsrv
+import gazebo_msgs.msg as gazmsg
+import gazebo_msgs.srv as gazsrv
+import geometry_msgs.msg import geomsg
+
+# import Gazebo-ROS related libraries
+from gazebo_ros import gazebo_interface
+
+# from gazebo_msgs.msg import *
+# from gazebo_msgs.srv import *
+# from geometry_msgs.msg import Point, Pose, Quaternion, Twist, Wrench, Vector3
+import tf.transformations as tft
+
+# import PRL
+from ros_rbdl import ROS_RBDL
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class GazeboROS(ROS_RBDL):
+ r"""Gazebo ROS RBDL Interface
+
+ This simulator uses Gazebo as the simulator, ROS to communicate with this simulator (to send and receive any
+ information related to the simulator and the objects inside of it like the robots), and RBDL to compute the
+ kinematics and dynamics of the robots. This class acts as the main bridge that connects what happens between
+ the simulator Gazebo and the PyRoboLearn framework.
+
+ Examples::
+ from pyrobolearn.simulators import GazeboROS
+
+ sim = GazeboROS(render=True)
+
+ References:
+ [1] ROS: http://www.ros.org/
+ [2] Gazebo: http://gazebosim.org/
+ [3] RBDL: https://rbdl.bitbucket.io/
+
+ Repositories:
+ * Xacro package: https://github.com/ros/xacro
+ * Gazebo ROS packages: https://github.com/ros-simulation/gazebo_ros_pkgs
+ * ROS control packages: https://github.com/ros-controls/ros_control
+ """
+
+ def __init__(self, render=True, ros_master_uri=11316, gazebo_master_uri=11345):
+ super(GazeboROS, self).__init__()
+
+ # Environment variable
+ self.env = os.environ.copy()
+
+ self.env["ROS_MASTER_URI"] = "http://localhost:" + str(ros_master_uri)
+ self.env["GAZEBO_MASTER_URI"] = "http://localhost:" + str(gazebo_master_uri)
+
+ # this is for the rospy methods such as: wait_for_service(), init_node(), ...
+ os.environ['ROS_MASTER_URI'] = self.env['ROS_MASTER_URI']
+
+ # create ROS core
+ # subprocess.Popen("roscore", env=self.env)
+ self.ros_proc = subprocess.Popen(["roscore", "-p", str(ros_master_uri)], env=self.env,
+ preexec_fn=os.setsid) # , shell=True)
+
+ # create Gazebo ROS
+ self.gzserver_proc = None
+ self.gzclient_proc = None
+
+ # Gazebo Services
+ self.reset_srv = rospy.ServiceProxy('/gazebo/reset_simulation', stdsrv.Empty)
+ self.pause_srv = rospy.ServiceProxy('/gazebo/pause_physics', stdsrv.Empty)
+ self.unpause_srv = rospy.ServiceProxy('/gazebo/unpause_physics', stdsrv.Empty)
+ self.get_physics_properties_srv = rospy.ServiceProxy('/gazebo/get_physics_properties', stdsrv.Empty)
+ self.set_physics_properties_srv = rospy.ServiceProxy('/gazebo/set_physics_properties',
+ gazsrv.SetPhysicsProperties)
+
+ # keep a list of bodies
+ self.bodies = []
+
+ # Simulators
+
+ def reset(self):
+ """
+ Reset the Gazebo simulation.
+ """
+ rospy.wait_for_service('/gazebo/reset_simulation')
+ try:
+ self.reset_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/reset_simulation service call failed")
+
+ def close(self):
+ """
+ Close everything
+ """
+ # delete Gazebo
+ if self.gzclient_proc is not None:
+ os.killpg(os.getpgid(self.gzclient_proc.pid), signal.SIGTERM)
+ if self.gzserver_proc is not None:
+ os.killpg(os.getpgid(self.gzserver_proc.pid), signal.SIGTERM)
+
+ # delete ROS
+ os.killpg(os.getpgid(self.ros_proc.pid), signal.SIGTERM)
+
+ def seed(self, seed=None):
+ """Set the given seed in the simulator."""
+ if seed is None:
+ return []
+ rospy.wait_for_service('/gazebo/set_seed')
+ try:
+ rospy.ServiceProxy('/gazebo/set_seed', SetSeedSrv)(seed)
+ except rospy.ServiceException, e:
+ print("/GazeboRosGym/set_seed service call failed")
+ return [seed]
+
+ def step(self, sleep_dt=0):
+ """Perform a step in the simulator, and sleep the specified time."""
+ self.unpause()
+ time.sleep(sleep_dt)
+ # TODO apply stuffs in simulator
+ self.pause()
+
+ def render(self, flag=True):
+ """Render the simulation."""
+ if flag:
+ if self.gzclient_proc is None:
+ pass
+ else:
+ if self.gzclient_proc is not None:
+ pass
+
+ def set_time_step(self, time_step):
+ """Set the time step in the simulator."""
+ set_physics_request = self.get_physics_properties()
+
+ # set time step
+ set_physics_request.time_step = time_step
+
+ # set the physics properties
+ rospy.wait_for_service('/gazebo/set_physics_properties')
+ try:
+ self.set_physics_properties_srv(set_physics_request)
+ except rospy.ServiceException, e:
+ print("/gazebo/reset_simulation service call failed")
+
+ def set_real_time(self):
+ """Enable real time in the simulator."""
+ self.unpause()
+
+ def pause(self):
+ """Pause the simulator if in real-time."""
+ rospy.wait_for_service('/gazebo/pause_physics')
+ try:
+ self.pause_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/pause_physics service call failed")
+
+ def unpause(self):
+ """Unpause the simulator if in real-time."""
+ rospy.wait_for_service('/gazebo/unpause_physics')
+ try:
+ self.unpause_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/unpause_physics service call failed")
+
+ def get_physics_properties(self):
+ """Get the physics engine parameters."""
+ rospy.wait_for_service('/gazebo/get_physics_properties')
+ try:
+ srv = self.get_physics_properties_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/reset_simulation service call failed")
+ return srv
+
+ def set_physics_properties(self, *args, **kwargs):
+ """Set the physics engine parameters."""
+ rospy.wait_for_service('/gazebo/set_physics_properties')
+ try:
+ self.set_physics_properties_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/reset_simulation service call failed")
+
+ def start_logging(self, *args, **kwargs):
+ """Start the logging."""
+ pass
+
+ def stop_logging(self, logger_id):
+ """Stop the logging."""
+ pass
+
+ def set_gravity(self, gravity=(0, 0, -9.81)):
+ """Set the gravity in the simulator."""
+ set_physics_request = self.get_physics_properties()
+
+ # set attributes
+ set_physics_request.gravity.x = gravity[0]
+ set_physics_request.gravity.y = gravity[1]
+ set_physics_request.gravity.z = gravity[2]
+
+ # set the physics properties
+ rospy.wait_for_service('/gazebo/set_physics_properties')
+ try:
+ self.set_physics_properties_srv(set_physics_request)
+ except rospy.ServiceException, e:
+ print("/gazebo/reset_simulation service call failed")
+
+ def save(self, on_disk=False):
+ """Save the state of the simulator."""
+ pass
+
+ def load(self, state):
+ """Load the simulator to a previous state."""
+ pass
+
+ def load_plugin(self, plugin):
+ """Load a certain plugin in the simulator."""
+ pass
+
+ def execute_plugin_commands(self, plugin_id, commands):
+ """Execute the commands on the specified plugin."""
+ pass
+
+ def unload_plugin(self, plugin_id):
+ """Unload the specified plugin from the simulator."""
+ pass
+
+ # loading URDFs, SDFs, MJCFs
+
+ def load_urdf(self, filename, position, orientation):
+ """Load a URDF file in the simulator."""
+ robot_namespace = rospy.get_namespace()
+ gazebo_namespace = "/gazebo"
+ reference_frame = ""
+ model_name = filename.split('/')[-1].split('.')[0] # assume filename='path/to/file(.xacro).urdf'
+
+ # if xacro file, use xacro.py with the list of arguments
+
+ # load file
+ f = open(filename, 'r')
+ model_xml = f.read()
+ if model_xml == "":
+ rospy.logerr("Error: file is empty %s", filename)
+ sys.exit(0)
+
+ # create initial pose
+ initial_pose = geomsg.Pose()
+ initial_pose.position.x = position[0]
+ initial_pose.position.y = position[1]
+ initial_pose.position.z = position[2]
+ q = geomsg.Quaternion()
+ q.x = orientation[0]
+ q.y = orientation[1]
+ q.z = orientation[2]
+ q.w = orientation[3]
+ initial_pose.orientation = q
+
+ success = gazebo_interface.spawn_urdf_model_client(model_name, model_xml, robot_namespace, initial_pose,
+ reference_frame, gazebo_namespace)
+ if not success:
+ raise ValueError("Could not load the given URDF in Gazebo.")
+
+ body_id = len(self.bodies)
+ self.bodies.append(model_name)
+ return body_id
+
+ def load_sdf(self, filename):
+ """Load a SDF file in the simulator."""
+ robot_namespace = rospy.get_namespace()
+ gazebo_namespace = "/gazebo"
+ reference_frame = ""
+ position = (0., 0., 0.)
+ orientation = (0., 0., 0., 1.)
+ model_name = filename.split('/')[-1].split('.')[-2] # assume filename='path/to/file.sdf'
+
+ # load file
+ f = open(filename, 'r')
+ model_xml = f.read()
+ if model_xml == "":
+ rospy.logerr("Error: file is empty %s", filename)
+ sys.exit(0)
+
+ # create initial pose
+ initial_pose = geomsg.Pose()
+ initial_pose.position.x = position[0]
+ initial_pose.position.y = position[1]
+ initial_pose.position.z = position[2]
+ q = geomsg.Quaternion()
+ q.x = orientation[0]
+ q.y = orientation[1]
+ q.z = orientation[2]
+ q.w = orientation[3]
+ initial_pose.orientation = q
+
+ success = gazebo_interface.spawn_sdf_model_client(model_name, model_xml, robot_namespace, initial_pose,
+ reference_frame, gazebo_namespace)
+ if not success:
+ raise ValueError("Could not load the given SDF in Gazebo.")
+
+ body_id = len(self.bodies)
+ self.bodies.append(model_name)
+ return body_id
+
+ def load_mjcf(self, filename):
+ """Load MJCF file."""
+ raise NotImplementedError("Loading a MJCF xml file in Gazebo is currently not possible.")
+
+
+
+
+class GazeboROSEnv(gazebo_env.GazeboEnv):
+ """
+ This class defines the Gazebo - OpenAI Gym interface.
+ The communication between the 2 systems is done using ROS.
+ """
+
+ def __init__(self, roslaunch_filename, package_name, ros_master_uri=11316, gazebo_master_uri=11345):
+
+ if roslaunch_filename is None:
+ raise ValueError("Expecting the roslaunch filename to be different from None")
+ if package_name is None:
+ raise ValueError("Expecting the package name to be different from None")
+
+ # Environment variable
+ self.env = os.environ.copy()
+
+ self.env["ROS_MASTER_URI"] = "http://localhost:" + str(ros_master_uri)
+ self.env["GAZEBO_MASTER_URI"] = "http://localhost:" + str(gazebo_master_uri)
+
+ # this is for the rospy methods such as: wait_for_service(), init_node(), ...
+ os.environ['ROS_MASTER_URI'] = self.env['ROS_MASTER_URI']
+
+ # Roscore and init node
+ #subprocess.Popen("roscore", env=self.env)
+ print('ROSCORE...')
+ self.ros_proc = subprocess.Popen(["roscore", "-p", str(ros_master_uri)], env=self.env, preexec_fn=os.setsid) #, shell=True)
+
+ rospy.wait_for_service('/rosout/get_loggers')
+ print('REGISTERING NODE...')
+ rospy.init_node('gym', anonymous=True)
+
+ # Roslaunch
+ print('ROSLAUNCH...')
+ print(package_name)
+ print(roslaunch_filename)
+ self.roslaunch_proc = subprocess.Popen(["roslaunch", package_name, roslaunch_filename, 'gui:=false', 'paused:=true'],
+ env=self.env,
+ preexec_fn=os.setsid)
+ #shell=True)
+ self.gzclient_pid = 0
+
+ rospy.wait_for_service('/gazebo/reset_simulation')
+ print('ROSLAUNCH DONE')
+
+ # Gazebo Services
+ self.reset_srv = rospy.ServiceProxy('/gazebo/reset_simulation', stdSrv.Empty)
+ self.pause_srv = rospy.ServiceProxy('/gazebo/pause_physics', stdSrv.Empty)
+ self.unpause_srv = rospy.ServiceProxy('/gazebo/unpause_physics', stdSrv.Empty)
+
+ def _seed(self, seed):
+ """
+ Set the seed in Gazebo using the new defined service.
+ """
+ if seed is None: return []
+ rospy.wait_for_service('/GazeboRosGym/set_seed')
+ try:
+ rospy.ServiceProxy('/GazeboRosGym/set_seed', SetSeed)(seed)
+ except rospy.ServiceException, e:
+ print("/GazeboRosGym/set_seed service call failed")
+ return [seed]
+
+ def reset_simulation(self):
+ """
+ Reset the Gazebo simulation.
+ """
+ rospy.wait_for_service('/gazebo/reset_simulation')
+ try:
+ self.reset_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/reset_simulation service call failed")
+
+ def pause_physics(self):
+ """
+ Pause the Gazebo physics engine.
+ """
+ rospy.wait_for_service('/gazebo/pause_physics')
+ try:
+ self.pause_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/pause_physics service call failed")
+
+ def unpause_physics(self):
+ """
+ Unpause the Gazebo physics engine.
+ """
+ rospy.wait_for_service('/gazebo/unpause_physics')
+ try:
+ self.unpause_srv()
+ except rospy.ServiceException, e:
+ print("/gazebo/unpause_physics service call failed")
+
+ def configure(self, *args, **kwargs):
+ """
+ Configure Gazebo with the given parameters.
+
+ Example (using the various rosservice for Gazebo):
+ - set the PID parameters
+ - set link/joint properties
+ - set link state
+ - set model configuration/state
+ - set physics properties (time step & update rate)
+ """
+ raise NotImplementedError("This function needs to be overwritten...")
+
+ def act(self, action):
+ """
+ Apply the action in the environment.
+ """
+ raise NotImplementedError("This function needs to be overwritten...")
+
+ def get_state(self):
+ """
+ Return the state.
+
+ Example:
+ The observation could be an image, while the state could be the position
+ (and velocity) of a target on the picture. The state is used to compute
+ the reward function.
+ """
+ raise NotImplementedError("This function needs to be overwritten...")
+
+ def get_obs(self):
+ """
+ Return the observation.
+
+ Example:
+ The observation could be an image, while the state could be the position
+ (and velocity) of a target on the picture. The state is used to compute
+ the reward function.
+ """
+ raise NotImplementedError("This function needs to be overwritten...")
+
+ def get_state_and_obs(self):
+ """
+ Return the state and observation.
+
+ Example:
+ The observation could be an image, while the state could be the position
+ (and velocity) of a target on the picture. The state is used to compute
+ the reward function.
+ """
+ return get_state(), get_obs()
+
+ def compute_reward(self, state, obs):
+ """
+ Compute and return the reward based on the state and on the observation.
+ It also returns a boolean value indicating if the task is over or not.
+ """
+ raise NotImplementedError("This function needs to be overwritten...")
+
+ def _step(self, action):
+ """
+ Run one timestep in the simulator.
+ """
+ self.unpause_physics()
+
+ self.act(action) # should this be before unpause_physics?
+ state, obs = self.get_state_and_obs()
+
+ self.pause_physics()
+
+ reward, done = self.compute_reward(state, obs)
+ return obs, reward, done, state
+
+ def _reset(self):
+ """
+ Reset the simulator.
+ """
+ self.reset_simulation()
+ self.unpause_physics()
+ obs = self.get_state()[1]
+ self.pause_physics()
+ return obs
diff --git a/pyrobolearn/simulators/gazebo.py b/pyrobolearn/simulators/gazebo.py
new file mode 100644
index 0000000..e66d422
--- /dev/null
+++ b/pyrobolearn/simulators/gazebo.py
@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+"""Define the Gazebo Simulator API.
+
+This is the main interface that communicates with the Gazebo simulator [1]. By defining this interface, it allows to
+decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+Gazebo. Note that this simulator does not use any ROS packages.
+
+Warnings: The use of this simulator necessitates Python wrappers for the Gazebo simulator [1]. Currently, none are
+provided, and thus the interface defined here is currently unusable.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] Gazebo: http://gazebosim.org/
+"""
+
+from simulator import Simulator
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class Gazebo(Simulator):
+ r"""Gazebo Simulator interface.
+
+ References:
+ [1] Gazebo: http://gazebosim.org/
+ """
+
+ def __init__(self, render=True):
+ super(Gazebo, self).__init__(render=render)
+ raise NotImplementedError
diff --git a/pyrobolearn/simulators/mujoco.py b/pyrobolearn/simulators/mujoco.py
new file mode 100644
index 0000000..448f79b
--- /dev/null
+++ b/pyrobolearn/simulators/mujoco.py
@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+"""Define the MuJoCo Simulator API.
+
+This is the main interface that communicates with the MuJoCo simulator [1]. By defining this interface, it allows to
+decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+MuJoCo.
+
+Warnings: The MuJoCo simulator requires a license in order to use it.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] MuJoCo: http://www.mujoco.org/
+ [2] MuJoCo Python: https://github.com/openai/mujoco-py
+ [3] DeepMind Control Suite: https://github.com/deepmind/dm_control/tree/master/dm_control/mujoco
+"""
+
+from simulator import Simulator
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class Mujoco(Simulator):
+ r"""Mujoco Simulator interface.
+
+ This is the main interface that communicates with the MuJoCo simulator [1]. By defining this interface, it allows
+ to decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+ MuJoCo.
+
+ Warnings: The MuJoCo simulator requires a license in order to use it.
+
+ References:
+ [1] MuJoCo: http://www.mujoco.org/
+ [2] MuJoCo Python: https://github.com/openai/mujoco-py
+ [3] DeepMind Control Suite: https://github.com/deepmind/dm_control/tree/master/dm_control/mujoco
+ """
+
+ def __init__(self, render=True):
+ super(Mujoco, self).__init__(render=render)
+ raise NotImplementedError
diff --git a/pyrobolearn/simulators/opensim.py b/pyrobolearn/simulators/opensim.py
new file mode 100644
index 0000000..20f2c9b
--- /dev/null
+++ b/pyrobolearn/simulators/opensim.py
@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+"""Define the OpenSim Simulator API.
+
+This is the main interface that communicates with the OpenSim simulator [1]. By defining this interface, it allows to
+decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+OpenSim.
+
+Warnings: This simulator only works for musculoskeletal models.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] OpenSim: https://opensim.stanford.edu/
+ [2] OpenSim Core: https://github.com/opensim-org/opensim-core
+ [3] OpenSim Reinforcement Learning: https://github.com/stanfordnmbl/osim-rl
+"""
+
+from simulator import Simulator
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class OpenSim(Simulator):
+ r"""OpenSim simulator
+
+ """
+
+ def __init__(self):
+ super(OpenSim, self).__init__()
+ raise NotImplementedError
diff --git a/pyrobolearn/simulators/ros.py b/pyrobolearn/simulators/ros.py
new file mode 100644
index 0000000..a38c031
--- /dev/null
+++ b/pyrobolearn/simulators/ros.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+"""Define the Bullet Simulator API.
+
+This is the main interface that communicates with the PyBullet simulator [1]. By defining this interface, it allows to
+decouple the PyRoboLearn framework from the simulator. It also converts some data types to the ones required by
+PyBullet. For instance, some methods in PyBullet do not accepts numpy arrays but only lists. The interface provided
+here makes the necessary conversions.
+
+The signature of each method defined here are inspired by [1,2] but in accordance with the PEP8 style guide [3].
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] PyBullet: https://pybullet.org
+ [2] PyBullet Quickstart Guide: https://docs.google.com/document/d/10sXEhzFRSnvFcl3XxNGhnD4N2SedqwdAvK3dsihxVUA
+ [3] PEP8: https://www.python.org/dev/peps/pep-0008/
+"""
+
+import rospy
+from simulator import Simulator
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class ROSModel(object):
+ r"""ROS Model
+
+ """
+
+ def __init__(self, filename):
+ self.urdf = filename
+ # get ros services and ros topics from URDF
+
+ # create
+ pass
+
+
+class ROS(Simulator):
+ r"""ROS Interface
+ """
+
+ def __init__(self):
+ super(ROS, self).__init__()
+ self.models = []
+
+ def load_urdf(self, filename, position=None, orientation=None):
+ # load URDF: get ros services and ros topics
+ model = ROSModel(filename)
+
+ # create id and add model to the list of models
+ idx = len(self.models)
+ self.models.append(model)
+
+ # return id
+ return idx
diff --git a/pyrobolearn/simulators/ros_rbdl.py b/pyrobolearn/simulators/ros_rbdl.py
new file mode 100644
index 0000000..e4720a0
--- /dev/null
+++ b/pyrobolearn/simulators/ros_rbdl.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+"""ROS-RBDL simulator
+
+This 'simulator' is not per se a simulator, it communicates with the real robots in the real world using ROS [1], and
+computes any necessary kinematic and dynamics information using the RBDL library [2].
+
+Specifically, this 'simulator' starts the `roscore` (if not already running), then loads robot urdf models and creates
+the necessary topics/services, and uses the rigid body dynamics library to compute kinematic and dynamic information
+about the model.
+
+Dependencies in PRL:
+* `pyrobolearn.simulators.simulator.Simulator`
+
+References:
+ [1] ROS: http://www.ros.org/
+ [2] RBDL: https://rbdl.bitbucket.io/
+"""
+
+import rospy
+import rbdl
+
+from simulator import Simulator
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class ROS_RBDL(Simulator):
+ r"""ROS-RBDL Interface.
+
+ References:
+ [1] ROS: http://www.ros.org/
+ [2] RBDL: https://rbdl.bitbucket.io/
+ [3] RBDL in Python: https://rbdl.bitbucket.io/dd/dee/_python_example.html
+ """
+
+ def __init__(self):
+ super(ROS_RBDL, self).__init__()
+
+ def step(self):
+ """Perform a step in the simulator."""
+ pass
+
+ def load_urdf(self, filename, position, orientation):
+ # load the model in rbdl
+ model = rbdl.loadModel(filename)
diff --git a/pyrobolearn/simulators/simulator.py b/pyrobolearn/simulators/simulator.py
new file mode 100644
index 0000000..cef859a
--- /dev/null
+++ b/pyrobolearn/simulators/simulator.py
@@ -0,0 +1,403 @@
+#!/usr/bin/env python
+"""Define the Simulator API.
+
+All the simulators inherit from the interface defined here. This acts as a bridge between the simulator and
+the PyRoboLearn framework. The signature of each method presents in this interface were inspired by the ones defined
+in PyBullet [1,2], but in accordance with the PEP8 style guide [3].
+
+Because the simulator is based on the PyBullet API and we want all the simulator APIs to be similar, all the other
+simulators would have to be able to carry out operations such as querying the state of the robots, kinematics and
+dynamics, .
+
+Dependencies in PRL: None
+
+References:
+ [1] PyBullet: https://pybullet.org
+ [2] PyBullet Quickstart Guide: https://docs.google.com/document/d/10sXEhzFRSnvFcl3XxNGhnD4N2SedqwdAvK3dsihxVUA
+ [3] PEP8: https://www.python.org/dev/peps/pep-0008/
+"""
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class Simulator(object):
+ r"""Simulator (abstract class)
+
+ All the simulators inherits from the Simulator defined here. This acts as a bridge between the simulator and
+ the PyRoboLearn framework. This avoids the PyRoboLearn framework to depends on a particular simulator.
+ The signature of each method presents in this interface were inspired by the ones defined in PyBullet [1].
+
+ Examples::
+ sim = Bullet()
+ sim = ROS_RBDL()
+ sim = GazeboROS()
+
+ References:
+ [1] PyBullet: https://pybullet.org
+ [2] PEP8: https://www.python.org/dev/peps/pep-0008/
+ """
+
+ def __init__(self, render=True):
+ self._render = render
+ self.real_time = False
+
+ ##############
+ # Properties #
+ ##############
+
+ @property
+ def version(self):
+ """Return the version of the simulator."""
+ return 0
+
+ #############
+ # Operators #
+ #############
+
+ def __repr__(self):
+ """Return a string about the class for debugging and development."""
+ return self.__class__.__name__
+
+ def __str__(self):
+ """Return a readable string about the class."""
+ return self.__class__.__name__
+
+ def __del__(self):
+ """Close/Delete the simulator."""
+ self.close()
+
+ ###########
+ # Methods #
+ ###########
+
+ # Simulators
+
+ def reset(self):
+ """Reset the simulator."""
+ pass
+
+ def close(self):
+ """Close the simulator."""
+ pass
+
+ def seed(self, seed=None):
+ """Set the given seed in the simulator."""
+ pass
+
+ def step(self, sleep_time=0):
+ """Perform a step in the simulator, and sleep the specified time."""
+ pass
+
+ def render(self, flag=True):
+ """Render the simulation."""
+ pass
+
+ def hide(self):
+ """Hide the GUI."""
+ self.render(False)
+
+ def set_time_step(self, time_step):
+ """Set the time step in the simulator."""
+ pass
+
+ def set_real_time(self):
+ """Enable real time in the simulator."""
+ pass
+
+ def pause(self):
+ """Pause the simulator if in real-time."""
+ pass
+
+ def unpause(self):
+ """Unpause the simulator if in real-time."""
+ pass
+
+ def get_physics_properties(self):
+ """Get the physics engine parameters."""
+ pass
+
+ def set_physics_properties(self, *args, **kwargs):
+ """Set the physics engine parameters."""
+ pass
+
+ def start_logging(self, *args, **kwargs):
+ """Start the logging."""
+ pass
+
+ def stop_logging(self, logger_id):
+ """Stop the logging."""
+ pass
+
+ def set_gravity(self, gravity=(0, 0, -9.81)):
+ """Set the gravity in the simulator."""
+ pass
+
+ def save(self, on_disk=False):
+ """Save the state of the simulator."""
+ pass
+
+ def load(self, state):
+ """Load the simulator to a previous state."""
+ pass
+
+ def load_plugin(self, plugin):
+ """Load a certain plugin in the simulator."""
+ pass
+
+ def execute_plugin_commands(self, plugin_id, commands):
+ """Execute the commands on the specified plugin."""
+ pass
+
+ def unload_plugin(self, plugin_id):
+ """Unload the specified plugin from the simulator."""
+ pass
+
+ # loading URDFs, SDFs, MJCFs
+
+ def load_urdf(self, filename, position, orientation):
+ """Load a URDF file in the simulator."""
+ pass
+
+ def load_sdf(self, filename):
+ """Load a SDF file in the simulator."""
+ pass
+
+ def load_mjcf(self, filename):
+ """Load a Mujoco file in the simulator."""
+ pass
+
+ def load_mesh(self, filename, position, orientation=(0, 0, 0, 1), mass=1., scale=(1., 1., 1.), color=(1, 1, 1, 1),
+ flags=None):
+ """Load a mesh into the simulator.
+
+ Args:
+ filename (str): path to file for the mesh. Currently, only Wavefront .obj. It will create convex hulls
+ for each object (marked as 'o') in the .obj file.
+ position (float[3]): position of the mesh in the Cartesian world space (in meters)
+ orientation (float[4], np.quaternion): orientation of the mesh using quaternion.
+ If np.quaternion then it uses the convention (w,x,y,z). If float[4], it uses the convention (x,y,z,w)
+ mass (float): mass of the mesh (in kg). If mass = 0, it won't move even if there is a collision.
+ scale (float[3]): scale the mesh in the (x,y,z) directions
+ color (int[4]): color of the mesh (by default: white and opaque)
+ flags (int, None): if flag = `sim.GEOM_FORCE_CONCAVE_TRIMESH` (=1), this will create a concave static
+ triangle mesh. This should not be used with dynamic/moving objects, only for static (mass=0) terrain.
+
+ Returns:
+ int: unique id of the mesh in the world
+ """
+ pass
+
+ # bodies
+
+ def create_visual_shape(self, shape_type, radius=0.5, half_extents=(1, 1, 1), length=1, filename='.obj'):
+ pass
+
+ def get_visual_shape_data(self, object_id):
+ pass
+
+ def create_collision_shape(self, shape_type, radius=0.5, half_extents=(1, 1, 1), length=1):
+ pass
+
+ def get_collision_shape_data(self):
+ pass
+
+ def create_body(self):
+ """Create a body in the simulator."""
+ pass
+
+ def remove_body(self, body_id):
+ """Remove a particular body in the simulator."""
+ pass
+
+ def num_bodies(self):
+ """Return the number of bodies present in the simulator."""
+ pass
+
+ def get_body_info(self, body_id):
+ """Get the specified body information."""
+ pass
+
+ def get_body_id(self):
+ pass
+
+ # constraint
+
+ def create_constraint(self):
+ pass
+
+ def remove_constraint(self):
+ pass
+
+ def change_constraint(self):
+ pass
+
+ def get_num_constraint(self):
+ pass
+
+ def get_constraint_id(self):
+ pass
+
+ def get_constraint_info(self):
+ pass
+
+ def get_constraint_state(self):
+ pass
+
+ # objects
+
+ def get_base_pose(self):
+ pass
+
+ def reset_base_pose(self):
+ pass
+
+ def get_base_position(self):
+ pass
+
+ def reset_base_position(self):
+ pass
+
+ def get_base_orientation(self):
+ pass
+
+ def reset_base_orientation(self):
+ pass
+
+ def get_base_velocity(self):
+ pass
+
+ def reset_base_velocity(self):
+ pass
+
+ def apply_external_force(self):
+ pass
+
+ def apply_external_torque(self):
+ pass
+
+ # robots (joints and links)
+
+ def get_num_joints(self):
+ pass
+
+ def get_joint_info(self):
+ pass
+
+ def get_joint_state(self):
+ pass
+
+ def get_joint_states(self):
+ pass
+
+ def reset_joint_state(self):
+ pass
+
+ def enable_joint_force_torque_sensor(self):
+ pass
+
+ def set_joint_motor_control(self):
+ pass
+
+ def set_joint_motor_control_array(self):
+ pass
+
+ def get_link_state(self):
+ pass
+
+ # visualization
+
+ def compute_view_matrix(self):
+ pass
+
+ def compute_projection_matrix(self):
+ pass
+
+ def get_camera_image(self):
+ pass
+
+ def load_texture(self):
+ pass
+
+ # collisions
+
+ def get_overlapping_objects(self):
+ pass
+
+ def get_aabb(self):
+ pass
+
+ def get_contact_points(self):
+ pass
+
+ def get_closest_points(self):
+ pass
+
+ def ray_test(self):
+ pass
+
+ def ray_test_batch(self):
+ pass
+
+ # kinematics and dynamics
+
+ def get_dynamics_info(self):
+ pass
+
+ def change_dynamics(self):
+ pass
+
+ def calculate_jacobian(self):
+ pass
+
+ def calculate_mass_matrix(self):
+ pass
+
+ def calculate_inverse_kinematics(self):
+ pass
+
+ def calculate_inverse_dynamics(self):
+ pass
+
+ def calculate_forward_dynamics(self):
+ pass
+
+ # debug
+
+ def add_user_debug_line(self):
+ pass
+
+ def add_user_debug_text(self):
+ pass
+
+ def add_user_debug_parameter(self):
+ pass
+
+ def add_user_data(self):
+ pass
+
+ def configure_debug_visualizer(self):
+ pass
+
+ def get_debug_visualizer(self):
+ pass
+
+ def reset_debug_visualizer(self):
+ pass
+
+ # events (mouse, keyboard)
+
+ def get_keyboard_events(self):
+ pass
+
+ def get_mouse_events(self):
+ pass
+
+ def get_mouse_and_keyboard_events(self):
+ pass
diff --git a/pyrobolearn/simulators/simureal.py b/pyrobolearn/simulators/simureal.py
new file mode 100644
index 0000000..9412a52
--- /dev/null
+++ b/pyrobolearn/simulators/simureal.py
@@ -0,0 +1,62 @@
+# This file defines an interface which is used by the robot classes.
+# This falls under the "Adapter" design pattern, where we add an abstraction
+# layer, by providing a common interface to different simulators and real robots.
+#
+# The UML diagram is depicted below:
+#
+# simuRealInterface -----------<> robot / gym-env
+# -----^-----
+# | |
+# ros_rbdl pybullet
+# |
+# ros_gazebo
+#
+# where the robot and gym-env classes only interact with children from env_interface.
+#
+# --- Example ---
+# env_gazebo = ros_gazebo()
+# robot = Robot(env_gazebo, 'path_to_urdf')
+# print(robot.getJointStates()) # will check the joint state in gazebo.
+# robot.drawCoM() # will draw a small sphere at the CoM in the gazebo simulator.
+#
+# env_bullet = pybullet()
+# robot.change_env(env_bullet) # change env and reload the urdf in the given env.
+# print(robot.getJointStates()) # will check the joint state in pybullet.
+# robot.drawCoM() # will draw a small sphere at the CoM in the pybullet simulator.
+#
+# env_ros = ros_rbdl() # assuming the real robot can send and recv msgs via
+# robot.change_env(env_ros) # rostopics/rosservices, you can interact with it.
+# print(robot.getJointStates()) # will check the joint state via ros.
+# robot.drawCoM() # return error as we can't draw in the real world.
+# ---------------
+#
+# You can thus interact with different simulators or the real robots.
+# Simulators: pybullet, pygazebo, ros-gazebo
+#
+# Warning: the name might change in the future.
+
+
+from abc import ABCMeta, abstractmethod
+
+
+class SimuRealInterface(object):
+ """Simulation-Reality Interface.
+ This abstract class must be inherited by any simulators, or real interfaces.
+ """
+ __metaclass__ = ABCMeta
+
+ def __init__(self):
+ pass
+
+ @abstractmethod
+ def stepSimulation(self):
+ raise NotImplementedError("Step simulation is not implemented.")
+
+ @abstractmethod
+ def render(self):
+ raise NotImplementedError()
+
+ @abstractmethod
+ def loadURDF(self, filename, position, orientation):
+ raise NotImplementedError()
+
diff --git a/pyrobolearn/utils/__init__.py b/pyrobolearn/utils/__init__.py
new file mode 100644
index 0000000..5efc039
--- /dev/null
+++ b/pyrobolearn/utils/__init__.py
@@ -0,0 +1,89 @@
+
+import inspect
+import types
+import numpy as np
+
+# Built-in functions
+
+def hasAttribute(object, name):
+ """Check if the given object has an attribute (variable or method) with the given name"""
+ return hasattr(object, name)
+
+def hasVariable(object, name):
+ """Check if the given object has a variable with the given name"""
+ attribute = getattr(object, name, None)
+ if attribute is not None:
+ if not callable(attribute):
+ return True
+ # if callable, it might be a callable object, a function, or method
+ # A variable can be an object or a function, but not a method.
+ return not isinstance(attribute, types.MethodType) # types.FunctionType
+ return False
+
+def hasMethod(object, name):
+ """Check if the given object has a method with the given name"""
+ method = getattr(object, name, None)
+ return inspect.ismethod(method)
+
+def isMethod(object):
+ """Check if the given object is a method"""
+ return inspect.ismethod(object)
+
+def isClass(object):
+ """Check if the given object is a class"""
+ return inspect.isclass(object)
+
+def isModule(object):
+ """Check if the given object is a module"""
+ return inspect.ismodule(object)
+
+def isList(object):
+ """Check if the given object is a list"""
+ return isinstance(object, list)
+
+def isTuple(object):
+ """Check if the given object is a tuple"""
+ return isinstance(object, tuple)
+
+def isNumpyArray(object):
+ """Check if the given object is a numpy array"""
+ return isinstance(object, np.ndarray)
+
+def isDict(object):
+ """Check if the given object is a dictionary"""
+ return isinstance(object, dict)
+
+def isSet(object):
+ """Check the given object is a set"""
+ return isinstance(object, set)
+
+def isNone(object):
+ """Check if the given object is None"""
+ return object is None
+
+def isInt(object):
+ """Check if the given object is an integer"""
+ return isinstance(object, int)
+
+def isFloat(object):
+ """Check if the given object is a float"""
+ return isinstance(object, float)
+
+def isStr(object):
+ """Check if the given object is a string"""
+ return isinstance(object, str)
+
+def isChar(object):
+ """Check if the given object is a character"""
+ if isinstance(object, str):
+ if len(object) == 1:
+ return True
+ return False
+
+def isBool(object):
+ """Check if the given object is a boolean"""
+ return isinstance(object, bool)
+
+def isComplex(object):
+ """Check if the given object is a complex number"""
+ return isinstance(object, complex)
\ No newline at end of file
diff --git a/pyrobolearn/utils/bullet_utils.py b/pyrobolearn/utils/bullet_utils.py
new file mode 100644
index 0000000..e41923e
--- /dev/null
+++ b/pyrobolearn/utils/bullet_utils.py
@@ -0,0 +1,172 @@
+# This file provides some utilities with the pybullet interface.
+
+
+class RGBColor(object):
+ red = (1, 0, 0)
+ green = (0, 1, 0)
+ blue = (0, 0, 1)
+ black = (0, 0, 0)
+ white = (1, 1, 1)
+ orange = (1, 0.647, 0)
+ dark_orange = (1, 0.549, 0)
+ yellow = (1, 1, 0)
+ pink = (1, 0.753, 0.796)
+ light_pink = (1, 0.714, 0.757)
+ deep_pink = (1, 0.078, 0.576)
+ grey = (0.502, 0.502, 0.502)
+
+
+class RGBAColor(object):
+ alpha = 1 # 0 = transparent, 1 = opaque
+ red = (1, 0, 0, alpha)
+ green = (0, 1, 0, alpha)
+ blue = (0, 0, 1, alpha)
+ black = (0, 0, 0, alpha)
+ white = (1, 1, 1, alpha)
+ orange = (1, 0.647, 0, alpha)
+ dark_orange = (1, 0.549, 0, alpha)
+ yellow = (1, 1, 0, alpha)
+ pink = (1, 0.753, 0.796, alpha)
+ light_pink = (1, 0.714, 0.757, alpha)
+ deep_pink = (1, 0.078, 0.576, alpha)
+ grey = (0.502, 0.502, 0.502, alpha)
+
+
+class Key(object): # BulletKeys
+ """Map keys to ascii and bullet id"""
+ a = 97
+ b = 98
+ c = 99
+ d = 100
+ e = 101
+ f = 102
+ g = 103
+ h = 104
+ i = 105
+ j = 106
+ k = 107
+ l = 108
+ m = 109
+ n = 110
+ o = 111
+ p = 112
+ q = 113
+ r = 114
+ s = 115
+ t = 116
+ u = 117
+ v = 118
+ w = 119
+ x = 120
+ y = 121
+ z = 122
+ n0 = 48
+ n1 = 49
+ n2 = 50
+ n3 = 51
+ n4 = 52
+ n5 = 53
+ n6 = 54
+ n7 = 55
+ n8 = 56
+ n9 = 57
+ space = 32
+ shift = 65306
+ ctrl = 65307
+ alt = 65308
+ enter = 65309
+ left_arrow = 65295
+ right_arrow = 65296
+ top_arrow = 65297
+ bottom_arrow = 65298
+
+ # state
+ nothing = 0
+ down = 1
+ triggered = 2
+ pressed = 3
+ released = 4
+
+ # def __init__(self):
+ # # add symbols (<,>,[,',...) and numbers (0,1,2,...)
+ # keys = {chr(i): i for i in range(32, 65)}
+ # # add letters and symbols
+ # keys.update({chr(i): i for i in range(91, 127)})
+ # keys.update({char: i for char, i in zip(['shift', 'ctrl', 'alt', 'enter'] + ['left','right','top','bottom'],
+ # list(range(65306, 65310)) + list(range(65295,65299)))})
+ #
+ # self.keys = keys
+ # self.keystr = {value: key for key, value in keys.items()}
+ #
+ # self.a = 97
+ # self.b = 98
+ # self.c = 99
+ # self.d = 100
+ # self.e = 101
+ # self.f = 102
+ # self.g = 103
+ # self.h = 104
+ # self.i = 105
+ # self.j = 106
+ # self.k = 107
+ # self.l = 108
+ # self.m = 109
+ # self.n = 110
+ # self.o = 111
+ # self.p = 112
+ # self.q = 113
+ # self.r = 114
+ # self.s = 115
+ # self.t = 116
+ # self.u = 117
+ # self.v = 118
+ # self.w = 119
+ # self.x = 120
+ # self.y = 121
+ # self.z = 122
+ # self.n0 = 48
+ # self.n1 = 49
+ # self.n2 = 50
+ # self.n3 = 51
+ # self.n4 = 52
+ # self.n5 = 53
+ # self.n6 = 54
+ # self.n7 = 55
+ # self.n8 = 56
+ # self.n9 = 57
+ # self.space = 32
+ # self.shift = 65306
+ # self.ctrl = 65307
+ # self.alt = 65308
+ # self.enter = 65309
+ # self.left_arrow = 65295
+ # self.right_arrow = 65296
+ # self.top_arrow = 65297
+ # self.bottom_arrow = 65298
+ #
+ # self.nothing = 0
+ # self.down = 1
+ # self.triggered = 2
+ # self.pressed = 3
+ # self.released = 4
+
+
+class Mouse(object): # Bullet mouse
+
+ # event type
+ moving = 1
+ button = 2
+
+ # button index
+ no_click = -1
+ left_click = 0
+ middle_click = 1 # scroll
+ right_click = 2
+
+ # button state
+ # state
+ nothing = 0
+ down = 1 # (never observed)
+ triggered = 2 # (never observed)
+ pressed = 3
+ released = 4
diff --git a/pyrobolearn/utils/cmu_mocap_parser.py b/pyrobolearn/utils/cmu_mocap_parser.py
new file mode 100644
index 0000000..e69de29
diff --git a/pyrobolearn/utils/converter.py b/pyrobolearn/utils/converter.py
new file mode 100644
index 0000000..762bf58
--- /dev/null
+++ b/pyrobolearn/utils/converter.py
@@ -0,0 +1,484 @@
+# This file describes converter classes which allows to convert from one certain data type to another.
+
+from abc import ABCMeta, abstractmethod
+import numpy as np
+import torch
+import quaternion
+import collections
+
+
+def roll(lst, shift):
+ """Roll elements of a list. This is similar to `np.roll()`"""
+ return lst[-shift:] + lst[:-shift]
+
+
+def numpy_to_torch(tensor):
+ return torch.from_numpy(tensor)
+
+
+def torch_to_numpy(tensor):
+ if tensor.requires_grad:
+ return tensor.detach().numpy()
+ return tensor.numpy()
+
+
+class TypeConverter(object):
+ r"""Type Converter class
+
+ It describes how to convert a type to another type, and inversely. For instance, a numpy array to a pytorch Tensor,
+ and vice-versa.
+ """
+ __metaclass__ = ABCMeta
+
+ def __init__(self, from_type, to_type):
+ self.from_type = from_type
+ self.to_type = to_type
+
+ @property
+ def from_type(self):
+ return self._from_type
+
+ @from_type.setter
+ def from_type(self, from_type):
+ if from_type is not None:
+ if isinstance(from_type, collections.Iterable):
+ for t in from_type:
+ if not isinstance(t, type):
+ raise TypeError("Expecting the from_type to be an instance of 'type'")
+ else:
+ if not isinstance(from_type, type):
+ raise TypeError("Expecting the from_type to be an instance of 'type'")
+ self._from_type = from_type
+
+ @property
+ def to_type(self):
+ return self._to_type
+
+ @to_type.setter
+ def to_type(self, to_type):
+ if to_type is not None:
+ if isinstance(to_type, collections.Iterable):
+ for t in to_type:
+ if not isinstance(t, type):
+ raise TypeError("Expecting the to_type to be an instance of 'type'")
+ else:
+ if not isinstance(to_type, type):
+ raise TypeError("Expecting the to_type to be an instance of 'type'")
+ self._to_type = to_type
+
+ @abstractmethod
+ def convertFrom(self, data):
+ """Convert to the 'from_type'"""
+ raise NotImplementedError
+
+ @abstractmethod
+ def convertTo(self, data):
+ """Convert to the 'to_type'"""
+ raise NotImplementedError
+
+ def convert(self, data):
+ """
+ Convert the data to the other type.
+ """
+ if isinstance(data, self.from_type): # or self.from_type is None:
+ return self.convertTo(data)
+ return self.convertFrom(data)
+
+ def __call__(self, data):
+ """
+ Call the convert method, and return the converted data.
+ """
+ return self.convert(data)
+
+
+class IdentityConverter(TypeConverter):
+ r"""Identity Converter
+
+ Dummy converter which does not convert the data.
+ """
+
+ def __init__(self):
+ super(IdentityConverter, self).__init__(None, None)
+
+ def convertFrom(self, data):
+ return data
+
+ def convertTo(self, data):
+ return data
+
+
+class NumpyListConverter(TypeConverter):
+ r"""Numpy - list converter
+
+ Convert lists/tuples to numpy arrays, and inversely.
+ """
+
+ def __init__(self, convention=0):
+ """Initialize the converter.
+
+ Args:
+ convention (int): convention to follow if 1D array. 0 to left it untouched, 1 to get column vector (i.e.
+ shape=(-1,1)), 2 to get row vector (i.e. shape=(1,-1)).
+ """
+ super(NumpyListConverter, self).__init__(from_type=(list, tuple), to_type=np.ndarray)
+
+ # check convention
+ if not isinstance(convention, int):
+ raise TypeError("Expecting an integer for the convention {0,1,2}")
+ if convention < 0 or convention > 2:
+ raise ValueError("Expecting the convention to belong to {0,1,2}")
+ self.convention = convention
+
+ def convertFrom(self, data):
+ """Convert to list"""
+ if isinstance(data, self.from_type):
+ return list(data)
+ elif isinstance(data, self.to_type):
+ if len(data.shape) == 2 and (data.shape[0] == 1 or data.shape[1] == 1):
+ return data.ravel().tolist() # flatten data
+ return data.tolist()
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to numpy array"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ data = np.array(data)
+ if len(data.shape) == 1:
+ if self.convention == 0: # left untouched
+ return data
+ elif self.convention == 1: # column vector
+ return data[:,np.newaxis]
+ else: # row vector
+ return data[np.newaxis,:]
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def reshape(self, data, shape):
+ """Reshape the data using the converter. Only valid if data is numpy array."""
+ if not isinstance(data, self.to_type):
+ data = self.convertTo(data)
+ return data.reshape(shape)
+
+ def transpose(self, data):
+ """Transpose the data using the converter"""
+ if not isinstance(data, self.to_type):
+ data = self.convertTo(data)
+ return data.T
+
+
+class QuaternionListConverter(TypeConverter):
+ r"""Quaternion - list converter
+
+ Convert a list/tuple to a quaternion, and vice-versa.
+ """
+
+ def __init__(self, convention=0):
+ """Initialize converter
+
+ Args:
+ convention (int): if 0, convert np.quaternion (w,x,y,z) to list [w,x,y,z], and inversely
+ if 1, convert np.quaternion (w,x,y,z) to list [x,y,z,w], and inversely
+ """
+ super(QuaternionListConverter, self).__init__(from_type=(list, tuple), to_type=np.quaternion)
+ if not isinstance(convention, int) or convention < 0 or convention > 1:
+ raise TypeError("Expecting convention to be 0 or 1.")
+ self.convention = convention
+
+ def convertFrom(self, data):
+ """Convert to list"""
+ if isinstance(data, self.from_type):
+ return list(data)
+ elif isinstance(data, self.to_type):
+ return np.roll(quaternion.as_float_array(data), -self.convention).tolist()
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to quaternion"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ return np.quaternion(*roll(data, -self.convention))
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+
+class QuaternionNumpyConverter(TypeConverter):
+ r"""Quaternion - numpy array converter
+
+ Convert a numpy array to a quaternion, and vice-versa.
+ """
+
+ def __init__(self, convention=0):
+ """Initialize converter
+
+ Args:
+ convention (int): if 0, convert np.quaternion (w,x,y,z) to list [w,x,y,z], and inversely
+ if 1, convert np.quaternion (w,x,y,z) to list [x,y,z,w], and inversely
+ """
+ super(QuaternionNumpyConverter, self).__init__(from_type=np.ndarray, to_type=np.quaternion)
+ if not isinstance(convention, int) or convention < 0 or convention > 1:
+ raise TypeError("Expecting convention to be 0 or 1.")
+ self.convention = convention
+
+ def convertFrom(self, data):
+ """Convert to numpy array"""
+ if isinstance(data, self.from_type):
+ return data
+ elif isinstance(data, self.to_type):
+ return np.roll(quaternion.as_float_array(data), -self.convention)
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to quaternion"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ return np.quaternion(roll(data.ravel().tolist(), -self.convention))
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def reshape(self, data, shape):
+ """Reshape the data using the converter. Only valid if data is numpy array."""
+ if not isinstance(data, self.from_type):
+ data = self.convertFrom(data)
+ return data.reshape(shape)
+
+ def transpose(self, data):
+ """Transpose the data using the converter"""
+ if not isinstance(data, self.from_type):
+ data = self.convertFrom(data)
+ return data.T
+
+
+class QuaternionPyTorchConverter(TypeConverter):
+ r"""Quaternion - pytorch tensor converter
+
+ Convert a pytorch tensor to a quaternion, and vice-versa. Currently, it converts it first to a numpy array and
+ then the other type.
+ """
+
+ def __init__(self, convention=0):
+ """Initialize converter
+
+ Args:
+ convention (int): if 0, convert np.quaternion (w,x,y,z) to list [w,x,y,z], and inversely
+ if 1, convert np.quaternion (w,x,y,z) to list [x,y,z,w], and inversely
+ """
+ super(QuaternionPyTorchConverter, self).__init__(from_type=torch.Tensor, to_type=np.quaternion)
+ if not isinstance(convention, int) or convention < 0 or convention > 1:
+ raise TypeError("Expecting convention to be 0 or 1.")
+ self.convention = convention
+
+ def convertFrom(self, data):
+ """Convert to pytorch tensor"""
+ if isinstance(data, self.from_type):
+ return data
+ elif isinstance(data, self.to_type):
+ return torch.from_numpy(np.roll(quaternion.as_float_array(data), -self.convention))
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to quaternion"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ return np.quaternion(roll(data.view(-1).data.tolist(), -self.convention))
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def reshape(self, data, shape):
+ """Reshape the data using the converter. Only valid if data is numpy array."""
+ if not isinstance(data, self.from_type):
+ data = self.convertFrom(data)
+ return data.view(shape)
+
+ def transpose(self, data):
+ """Transpose the data using the converter"""
+ if not isinstance(data, self.from_type):
+ data = self.convertFrom(data)
+ return data.t()
+
+
+class NumpyNumberConverter(TypeConverter):
+ r"""Numpy - number Converter
+
+ Convert a number to a numpy array of dimension 0 or 1, and vice-versa.
+ """
+
+ def __init__(self, dim_array=1):
+ super(NumpyNumberConverter, self).__init__(from_type=(int, float), to_type=np.ndarray)
+
+ # dimension array
+ if not isinstance(dim_array, int):
+ raise TypeError("The 'dim_array' argument should be an integer.")
+ if dim_array < 0 or dim_array > 1:
+ raise ValueError("The 'dim_array' argument should be 0 or 1.")
+ self.dim_array = dim_array
+
+ def convertFrom(self, data):
+ """Convert to a number"""
+ if isinstance(data, self.from_type):
+ return data
+ elif isinstance(data, self.to_type):
+ dim = len(data.shape)
+ if dim == 0:
+ return data[()]
+ elif dim == 1:
+ return data[0]
+ else:
+ raise ValueError("The numpy array should have a shape length of 0 or 1.")
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to numpy array"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ if self.dim_array == 0:
+ return np.array(data)
+ return np.array([data])
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+
+class PyTorchListConverter(TypeConverter):
+ r"""Pytorch - list converter
+
+ Convert lists/tuples to pytorch tensors. Currently, it converts it first to a numpy array and then the other type.
+ """
+
+ def __init__(self, convention=0):
+ """Initialize the converter.
+
+ Args:
+ convention (int): convention to follow if 1D array. 0 to left it untouched, 1 to get column vector (i.e.
+ shape=(-1,1)), 2 to get row vector (i.e. shape=(1,-1)).
+ """
+ super(PyTorchListConverter, self).__init__(from_type=(tuple, list), to_type=torch.Tensor)
+
+ # check convention
+ if not isinstance(convention, int):
+ raise TypeError("Expecting an integer for the convention {0,1,2}")
+ if convention < 0 or convention > 2:
+ raise ValueError("Expecting the convention to belong to {0,1,2}")
+ self.convention = convention
+
+ def convertFrom(self, data):
+ """Convert to list"""
+ if isinstance(data, self.from_type):
+ return list(data)
+ elif isinstance(data, self.to_type):
+ data = data.numpy() # convert to numpy first
+ if len(data.shape) == 2 and (data.shape[0] == 1 or data.shape[1] == 1):
+ return data.ravel().tolist() # flatten data
+ return data.tolist()
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to pytorch tensor"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ data = np.array(data)
+ if len(data.shape) == 1:
+ if self.convention == 1: # column vector
+ data = data[:,np.newaxis]
+ elif self.convention == 2: # row vector
+ data = data[np.newaxis,:]
+ return torch.from_numpy(data)
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def reshape(self, data, shape):
+ """Reshape the data using the converter. Only valid if data is numpy array."""
+ if not isinstance(data, self.to_type):
+ data = self.convertTo(data)
+ return data.view(shape)
+
+ def transpose(self, data):
+ """Transpose the data using the converter"""
+ if not isinstance(data, self.to_type):
+ data = self.convertTo(data)
+ return data.t()
+
+
+class PyTorchNumpyConverter(TypeConverter):
+ r"""PyTorch - Numpy Converter
+
+ Convert numpy arrays to a pytorch tensors, and vice-versa.
+ """
+
+ def __init__(self):
+ super(PyTorchNumpyConverter, self).__init__(from_type=np.ndarray, to_type=torch.Tensor)
+
+ def convertFrom(self, data):
+ """Convert to numpy array"""
+ if isinstance(data, self.from_type):
+ return data
+ elif isinstance(data, self.to_type):
+ if data.requires_grad:
+ return data.detach().numpy()
+ return data.numpy()
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def convertTo(self, data):
+ """Convert to pytorch tensor"""
+ if isinstance(data, self.to_type):
+ return data
+ elif isinstance(data, self.from_type):
+ return torch.from_numpy(data)
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def reshape(self, data, shape):
+ """Reshape the data based on the type using the converter."""
+ if isinstance(data, self.from_type): # np
+ return data.reshape(shape)
+ elif isinstance(data, self.to_type): # torch
+ return data.view(shape)
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+ def transpose(self, data):
+ """Transpose the data using the converter"""
+ if isinstance(data, self.from_type): # np
+ return data.T
+ elif isinstance(data, self.to_type): # torch
+ return data.t()
+ else:
+ raise TypeError("Type not known: {}".format(type(data)))
+
+
+# class OpenCVNumpyConverter(TypeConverter):
+# pass
+
+
+if __name__ == '__main__':
+ converter = NumpyListConverter()
+ print("Using {}".format(converter.__class__.__name__))
+ a = np.array(range(4))
+ print("on np.array: a={} with type {}".format(a, type(a)))
+ b = converter(a)
+ print("converter(a) gives: {} with type {}".format(b, type(b)))
+ b = converter.convertFrom(a)
+ print("converter.convertFrom(a) gives: {} with type {}".format(b, type(b)))
+ b = converter.convertTo(a)
+ print("converter.convertTo(a) gives: {} with type {}".format(b, type(b)))
+
+ A = np.array(range(4)).reshape(2, 2)
+ print("on numpy matrix: \nA={} with type {}".format(A, type(A)))
+ b = converter(A)
+ print("converter(a) gives: {} with type {}".format(b, type(b)))
+ b = converter.convertFrom(A)
+ print("converter.convertFrom(a) gives: {} with type {}".format(b, type(b)))
+ b = converter.convertTo(A)
+ print("converter.convertTo(a) gives: \n{} with type {}".format(b, type(b)))
diff --git a/pyrobolearn/utils/data_structures/__init__.py b/pyrobolearn/utils/data_structures/__init__.py
new file mode 100644
index 0000000..d876ab7
--- /dev/null
+++ b/pyrobolearn/utils/data_structures/__init__.py
@@ -0,0 +1,8 @@
+
+# Define common data structures
+
+# Ordered sets
+from orderedset import *
+
+# Graph
+from graph import *
\ No newline at end of file
diff --git a/pyrobolearn/utils/data_structures/graph.py b/pyrobolearn/utils/data_structures/graph.py
new file mode 100644
index 0000000..96f47ad
--- /dev/null
+++ b/pyrobolearn/utils/data_structures/graph.py
@@ -0,0 +1,51 @@
+
+
+class Graph(object):
+ r"""Graph.
+
+ This class described the graph data structure.
+ """
+ pass
+
+
+class DirectedGraph(Graph):
+ r"""Directed Graph.
+
+ This class described the directed graph data structure.
+
+ The graph is described using a dictionary.
+ graph = {node: [[parent nodes], [child nodes]]}
+ """
+ class Root(object):
+ pass
+
+ def __init__(self):
+ self._root = self.Root()
+ self._graph = {self._root: []}
+
+ def addNode(self, parent, node):
+ """
+ Add a node to the graph.
+ """
+ pass
+
+ def getParents(self, node):
+ """
+ Return the parent nodes of the given node.
+ """
+ pass
+
+ def getChildren(self, node):
+ """
+ Return the child nodes of the given node.
+ """
+ pass
+
+
+class DirectedAcyclicGraph(DirectedGraph):
+ """Directed Acyclic Graph.
+
+ In this data structure, cycles are not allowed.
+ """
+ def __init__(self):
+ super(DirectedAcyclicGraph, self).__init__()
diff --git a/pyrobolearn/utils/data_structures/orderedset.py b/pyrobolearn/utils/data_structures/orderedset.py
new file mode 100644
index 0000000..039269a
--- /dev/null
+++ b/pyrobolearn/utils/data_structures/orderedset.py
@@ -0,0 +1,957 @@
+
+import collections
+
+class OrderedSet(collections.MutableSet):
+ r"""Ordered Set
+
+ This is my own implementation of an ordered set, and was inspired a bit from [1] and [2].
+ In this class, we internally use a set and an (ordered) list to describe an ordered set.
+
+ In this implementation, the `delete/remove/discard item`, `move item`, `insert item`, `pop item` (except last
+ item) operations are pretty expensive with a time complexity of O(N). However, operations such as get item` and
+ `set item` have a time complexity of O(1).
+
+ If you need to easily add items and get access to them, and don't need to remove, insert, and move items,
+ use this class.
+
+ Here are the time complexities for the average (and worst) case scenario (more info on [3,4]):
+ * Iterate: O(N)
+ * Copy: O(N)
+ * Get Length: O(1)
+ * Item in set: O(1) (worst: O(N))
+ * Subset in set (without respecting the order): O(K) (worst: O(N))
+ * Add/append item: O(1)
+ * Delete/remove/discard item: O(N)
+ * Pop last: O(1)
+ * Pop first: O(N)
+ * Pop given index: O(N)
+ * Insert item: O(N)
+ * Move item: O(N)
+ * Get item from key: O(1)
+ * Get items from slice: O(K)
+ * Set item from key: O(1) if it doesn't have to move the data, otherwise O(N)
+ * Set items from slice: O(K+N) (N because it might have to move some data to accomodate for the new items, see [4])
+ * Delete item from key: O(N)
+
+ * Is superset/subset (while respecting the order): O(N)
+ * Union:
+ * Intersection:
+ * Difference:
+
+ References:
+ [1] http://code.activestate.com/recipes/576694-orderedset/
+ [2] https://stackoverflow.com/questions/1653970/does-python-have-an-ordered-set
+ [3] http://bigocheatsheet.com/
+ [4] https://wiki.python.org/moin/TimeComplexity?
+ """
+
+ def __init__(self, iterator=None):
+ """
+ Initialize the ordered set, which basically contains a set and a list.
+
+ Args:
+ iterator: an iterator
+ """
+ self._set = set()
+ self._list = []
+
+ if isinstance(iterator, collections.Iterable):
+ for item in iterator:
+ self.add(item)
+
+ def add(self, item):
+ """
+ Add/Append an item to the ordered set.
+ Time complexity: O(1)
+ """
+ if item not in self._set:
+ self._set.add(item)
+ self._list.append(item)
+
+ # alias
+ append = add
+
+ def extend(self, iterator):
+ """
+ Extend the ordered set by adding/appending elements from the given iterator.
+ Time complexity: O(K) where K is the size of the iterator
+ """
+ for item in iterator:
+ self.add(item)
+
+ def insert(self, idx, item):
+ """
+ Insert an item into the set at the specified index. If the item is already in the set, it moves it
+ to the specified index.
+ Time complexity: O(N)
+ """
+ # check idx
+ idx = self._checkIndex(idx)
+ if item in self._set:
+ # move the item at the specified location
+ self.move(idx, item) # O(N)
+ else:
+ # add it
+ self._list.insert(idx, item) # O(N)
+ self._set.add(item) # O(1)
+
+
+ def move(self, idx, item):
+ """
+ Move an item to the specified index. If the item is not in the set, it raises a KeyError.
+ Time complexity: O(N)
+ """
+ # remove the item from the list/set
+ self.remove(item) # O(N)
+
+ # insert item
+ self.insert(idx, item) # O(N)
+
+
+ def discard(self, item):
+ """
+ Remove an item from the ordered set if it is a member. If the item is not a member do nothing.
+ Time complexity: O(N)
+ """
+ if item in self._set:
+ self._list.remove(item) # O(N)
+ self._set.remove(item) # O(1)
+
+ def remove(self, item):
+ """
+ Remove an item from the ordered set. If the item is not a member, it raises a KeyError.
+ Time complexity: O(N)
+ """
+ if item not in self._set:
+ raise KeyError(item)
+ self.discard(item)
+
+ def pop(self, index=None):
+ """
+ Remove and return an item of the ordered set at the specified index (default last).
+ Time complexity: O(1) if last, O(N) if first.
+ Args:
+ index: index in the ordered set.
+ """
+ if index is None: index = len(self._list)
+ index = self._checkIndex(index) # to be sure the index is valid
+ self._set.remove(self._list[index]) # O(1)
+ item = self._list.pop(index) # O(1) if last, O(N) if first
+ return item
+
+ def copy(self):
+ """
+ Return a shallow copy of an ordered set.
+ Time complexity: O(N)
+ """
+ return self.__class__(self)
+
+ def _checkIndex(self, idx):
+ """
+ Check the given index; if it is in the range of the ordered set, and if it is negative return the
+ corresponding positive index.
+ """
+ if not isinstance(idx, int):
+ raise TypeError("idx should be an integer.")
+ if idx > len(self._list) or idx < -len(self._list):
+ return KeyError(idx)
+ if idx < 0:
+ idx = len(self._list) + idx
+ return idx
+
+ def union(self, *others):
+ """
+ Return the union of sets as a new set.
+ """
+ s = self.copy()
+ s.update(*others)
+ return s
+
+ def update(self, *others):
+ """
+ Update a set with the union of itself and others.
+ """
+ for other in others:
+ self |= other
+
+ def intersection(self, *others):
+ """
+ Return the intersection of two or more sets as a new set.
+ """
+ s = self.copy()
+ s.intersection_update(*others)
+ return s
+
+ def intersection_update(self, *others):
+ """
+ Update a set with the intersection of itself and another.
+ """
+ for other in others:
+ self &= other
+
+ def difference(self, *others):
+ """
+ Return the difference of two or more sets as a new set; i.e. all elements that are in this set but not
+ the others.
+ """
+ s = self.copy()
+ s.difference_update(*others)
+ return s
+
+ def difference_update(self, *others):
+ """
+ Remove all elements of another set from this set.
+ """
+ for other in others:
+ self -= other
+
+ def symmetric_difference(self, *others):
+ """
+ Return the symmetric difference of several sets as a new set; i.e. union of sets - intersection of sets.
+ """
+ s = self.copy()
+ s.symmetric_difference_update(*others)
+ return s
+
+ def symmetric_difference_update(self, *others):
+ """
+ Update a set with the symmetric difference of itself and others; i.e. set = union(sets) - intersection(sets)
+ """
+ intersection = self.copy()
+ self.update(*others) # compute union
+ intersection.intersection_update(*others) # compute intersection
+ self -= intersection
+
+ def issuperset(self, other, order=True):
+ """
+ Return True if the other set is a subset of this set. If 'order' is True, then the other set has to be
+ a subset of this set, and have the same order as this one.
+ Time complexity: O(N) if order, O(K) otherwise where N is the size of this set, and K is the size of the
+ other set.
+ """
+ if not isinstance(other, (OrderedSet, set)):
+ raise TypeError("The 'other' argument should be a set, or an ordered set.")
+ if len(other) == 0: # the empty set is always a subset of a set
+ return True
+ if len(other) > len(self): # the other set is bigger than this set, and thus is not a subset of that one
+ return False
+
+ # take into account the order if specified
+ if order:
+ if not isinstance(other, OrderedSet):
+ raise TypeError("The 'other' argument should be an ordered set.")
+
+ # traverse the subset and check each element appeared in the same order in the set
+ iterator = iter(self._list)
+ for item in other:
+ # check if item of subset is in the set
+ if item not in self:
+ return False
+
+ # traverse the ordered set until we find this item
+ while True:
+ try:
+ curr = next(iterator)
+ if curr == item: break
+ except StopIteration:
+ return False
+
+ # return True as we checked that all the items in the subset are in the set
+ return True
+
+ # the order is not important
+ else:
+ return all([(item in self) for item in other])
+
+ # alias
+ contains = issuperset
+
+ def issubset(self, other, order=True):
+ """
+ Return True if the other set is a superset of this set; i.e. return true if this set is a subset of the
+ other set.
+ Time complexity: same as `issuperset()`.
+ """
+ return other.issuperset(self, order=order)
+
+ def __repr__(self):
+ return '%s(%r)' % (self.__class__.__name__, list(self))
+
+ def __contains__(self, item):
+ """
+ Check if the given item/subset is in the set. The subset doesn't have the same order.
+ If the order is important, see the `issubset()` method.
+ Time complexity: O(1) if single item, O(K) if subset (where K is the size of the subset)
+ """
+ if isinstance(item, OrderedSet):
+ return self.issuperset(item, order=False)
+ return item in self._set
+
+ def __len__(self):
+ """
+ Return the length of the ordered set.
+ Time complexity: O(1)
+ """
+ return len(self._list)
+
+ def __iter__(self):
+ """
+ Iterate over the ordered set in the order the items have been added.
+ Time complexity: O(N)
+ """
+ for item in self._list:
+ yield item
+
+ def __reversed__(self):
+ """
+ Iterate over the ordered set in the reverse order the items have been added.
+ Time complexity: O(N)
+ """
+ for item in reversed(self._list):
+ yield item
+
+ def __getitem__(self, idx):
+ """
+ Return the item (ordered set) associated to the given index (indices).
+ Time complexity: O(1) if index, O(K) if slice (where K is the size of slice)
+ """
+ # checks
+ if len(self) == 0:
+ raise KeyError("Trying to get an item from an empty set.")
+ if isinstance(idx, int):
+ return self._list[idx]
+ else: # slice
+ return OrderedSet(self._list[idx])
+
+ def __setitem__(self, idx, item):
+ """
+ Replace the item at the specified index. If the item is already in the ordered set, it will move it to the
+ specified index.
+ Time complexity: O(1) if index is an int and it doesn't have to move an item, O(N) if it has to move it,
+ and O(K+N) if slice
+ """
+ if isinstance(idx, int):
+ if item in self._set:
+ self.move(idx, item)
+ else:
+ self._list[idx] = item
+ self._set.add(item)
+ elif isinstance(idx, slice): # slice
+ # replace in list
+ items_to_remove = self._list[idx] # O(K)
+ self._list[idx] = item # O(K+N)
+
+ # remove previous items from the set
+ for elem in items_to_remove:
+ self._set.remove(elem) # O(1)
+
+ # add new items in the set
+ for elem in item:
+ if elem not in self._set:
+ self._set.add(elem)
+ else:
+ raise TypeError("Expecting idx to be an int or slice.")
+
+ def __add__(self, other):
+ """
+ Union between two ordered sets.
+ """
+ return self | other
+
+ def __iadd__(self, other):
+ """
+ Update a set with the union of itself and the other set.
+ """
+ self |= other
+
+ def __and__(self, other):
+ """
+ Intersection between two ordered sets.
+ """
+ return super(OrderedSet, other).__and__(self)
+
+ def __iand__(self, other):
+ """
+ Update a set with the intersection of itself and the other set.
+ """
+ super(OrderedSet, self).__iand__(other)
+
+ def __mul__(self, other):
+ """
+ Intersection between two ordered sets.
+ """
+ return self & other
+
+ def __imul__(self, other):
+ """
+ Update a set with the intersection of itself and the other set.
+ """
+ self &= other
+
+
+################################################################################################
+
+
+class OrderedSet2(collections.MutableSet):
+ r"""Ordered Set
+
+ This is my own implementation of an ordered set, and was inspired a bit from [1] and [2].
+ In this class, we internally use a dictionary where keys are the items of the ordered set,
+ and each associated value is a tuple containing the pointer to the previous and next items.
+ It can thus be seen as a double-linked list with fast access.
+
+ In this implementation, the `get item`, `set item`, `move item`, and `insert item` operations are pretty
+ expensive with a time complexity of O(N) compared to a list (which has O(1)). However, operations such as
+ `delete/remove/discard item`, and `pop first/last items` have a time complexity of O(1).
+
+ If you need to easily remove and append items, and you don't need to access (get/set) the items in the set,
+ use this class.
+
+ Here are the time complexities for the average (and worst) case scenario (more info on [3,4]):
+ * Iterate: O(N)
+ * Copy: O(N)
+ * Get Length: O(1)
+ * Item in set: O(1) (worst: O(N))
+ * Subset in Set (without respecting the order): O(K) (worst: O(N))
+ * Add/append item: O(1)
+ * Delete/remove/discard item: O(1)
+ * Pop last: O(1)
+ * Pop first: O(1)
+ * Pop given index: O(N)
+ * Insert item: O(N)
+ * Move item: O(N)
+ * Get item from key: O(N)
+ * Get items from slice: O(N+K) (+K because we build a new ordered set)
+ * Set item from key: O(N)
+ * Set items from slice: Not Implemented
+ * Delete item from key: O(N)
+
+ * Is superset/subset (while respecting the order): O(N)
+ * Union:
+ * Intersection:
+ * Difference:
+
+ References:
+ [1] http://code.activestate.com/recipes/576694-orderedset/
+ [2] https://stackoverflow.com/questions/1653970/does-python-have-an-ordered-set
+ [3] http://bigocheatsheet.com/
+ [4] https://wiki.python.org/moin/TimeComplexity?
+ """
+
+ class _NonePtr(object): pass
+ NonePtr = _NonePtr()
+
+ def __init__(self, iterator=None):
+ """
+ Initialize the ordered set.
+
+ Args:
+ iterator: An iterator
+ """
+ self._start, self._end = self.NonePtr, self.NonePtr
+ self._map = {}
+
+ if isinstance(iterator, collections.Iterable):
+ for item in iterator:
+ self.add(item)
+
+ def add(self, item):
+ """
+ Add/Append an item to the ordered set.
+ Time complexity: O(1)
+ """
+ if item not in self._map:
+ if self._end == self.NonePtr: # first item
+ self._map[item] = [self.NonePtr, self.NonePtr]
+ self._start = item
+ self._end = item
+ else: # subsequent item
+ # update previous item to point to new item
+ self._map[self._end][1] = item
+ # append new item at the end
+ self._map[item] = [self._end, self.NonePtr]
+ # update end pointer
+ self._end = item
+
+ # alias
+ append = add
+
+ def extend(self, iterator):
+ """
+ Extend the ordered set by adding/appending elements from the given iterator.
+ Time complexity: O(K) where K is the size of the iterator
+ """
+ for item in iterator:
+ self.add(item)
+
+ def insert(self, idx, item):
+ """
+ Insert an item into the set at the specified index. If the item is already in the set, it moves it
+ to the specified index.
+ Time complexity: O(N)
+ """
+ # check idx
+ idx = self._checkIndex(idx)
+
+ # if the set is initially empty or index is the size of the set, just add the item (at the end)
+ if len(self._map) == 0 or idx == len(self._map):
+ self.append(item)
+ else:
+ if item in self._map:
+ self.move(idx, item)
+ else:
+ # get current item at the specified index, update the items nearby, and insert the new item
+ curr = self[idx]
+ prev_item, next_item = self._map[curr]
+ if prev_item == self.NonePtr: # beginning of the ordered set (idx=0)
+ self._map[curr][0] = item
+ self._map[item] = [self.NonePtr, curr]
+ else: # somewhere between the start and the end (not included)
+ self._map[item] = [prev_item, next_item]
+ self._map[prev_item][1] = item
+ self._map[next_item][0] = item
+
+ def move(self, idx, item):
+ """
+ Move an item to the specified index. If the item is not in the set, it raises a ValueError.
+ Time complexity: O(N)
+ """
+ if item not in self._map:
+ return ValueError("The given item is not in the set.")
+
+ # remove the item from the list/set (time complexity: O(1))
+ self.remove(item)
+
+ # insert item
+ self.insert(idx, item)
+
+
+ def discard(self, item):
+ """
+ Remove an item from the ordered set if it is a member. If the item is not a member do nothing.
+ Time complexity: O(1)
+ """
+ if item in self._map:
+ prev_item, next_item = self._map[item]
+
+ # update previous item
+ if prev_item != self.NonePtr:
+ self._map[prev_item][1] = next_item
+ else: # we are removing the first item
+ self._start = next_item
+
+ # update next item
+ if next_item != self.NonePtr:
+ self._map[next_item][0] = prev_item
+ else: # we are removing the last item
+ self._end = prev_item
+
+ # remove item
+ self._map.pop(item)
+
+ def remove(self, item):
+ """
+ Remove an item from the ordered set. If the item is not a member, it raises a KeyError.
+ Time complexity: O(1)
+ """
+ if item not in self._map:
+ raise KeyError(item)
+ self.discard(item)
+
+ def pop(self, last=True):
+ """
+ Remove and return the first or last element of the ordered set depending on the provided argument.
+ Time complexity: O(1)
+ Args:
+ last: if True, remove and return the last item added to the set. If False, remove and return the first one.
+ """
+ if last:
+ item = self._end
+ else:
+ item = self._start
+ self.remove(item)
+ return item
+
+ def copy(self):
+ """
+ Return a shallow copy of an ordered set.
+ Time complexity: O(N)
+ """
+ return self.__class__(self)
+
+ def _checkIndex(self, idx):
+ """
+ Check the given index; if it is in the range of the ordered set, and if it is negative return the
+ corresponding positive index.
+ """
+ if not isinstance(idx, int):
+ raise TypeError("idx should be an integer.")
+ if idx > len(self._map) or idx < -len(self._map):
+ return KeyError(idx)
+ if idx < 0:
+ idx = len(self._map) + idx
+ return idx
+
+ def union(self, *others):
+ """
+ Return the union of sets as a new set.
+ """
+ s = self.copy()
+ s.update(*others)
+ return s
+
+ def update(self, *others):
+ """
+ Update a set with the union of itself and others.
+ """
+ for other in others:
+ self |= other
+
+ def intersection(self, *others):
+ """
+ Return the intersection of two or more sets as a new set.
+ """
+ s = self.copy()
+ s.intersection_update(*others)
+ return s
+
+ def intersection_update(self, *others):
+ """
+ Update a set with the intersection of itself and another.
+ """
+ for other in others:
+ self &= other
+
+ def difference(self, *others):
+ """
+ Return the difference of two or more sets as a new set; i.e. all elements that are in this set but not
+ the others.
+ """
+ s = self.copy()
+ s.difference_update(*others)
+ return s
+
+ def difference_update(self, *others):
+ """
+ Remove all elements of another set from this set.
+ """
+ for other in others:
+ self -= other
+
+ def symmetric_difference(self, *others):
+ """
+ Return the symmetric difference of several sets as a new set; i.e. union of sets - intersection of sets.
+ """
+ s = self.copy()
+ s.symmetric_difference_update(*others)
+ return s
+
+ def symmetric_difference_update(self, *others):
+ """
+ Update a set with the symmetric difference of itself and others; i.e. set = union(sets) - intersection(sets)
+ """
+ intersection = self.copy()
+ self.update(*others) # compute union
+ intersection.intersection_update(*others) # compute intersection
+ self -= intersection
+
+ def issuperset(self, other, order=True):
+ """
+ Return True if the other set is a subset of this set. If 'order' is True, then the other set has to be
+ a subset of this set, and have the same order as this one.
+ Time complexity: O(N) if order, O(K) otherwise where N is the size of this set, and K is the size of the
+ other set.
+ """
+ if not isinstance(other, (OrderedSet, set)):
+ raise TypeError("The 'other' argument should be a set or an ordered set.")
+ if len(other) == 0: # the empty set is always a subset of a set
+ return True
+ if len(other) > len(self): # the other set is bigger than this set, and thus is not a subset of that one
+ return False
+
+ # take into account the order if specified
+ if order:
+ if not isinstance(other, OrderedSet):
+ raise TypeError("The 'other' argument should be an ordered set.")
+
+ # check first item in the subset
+ curr_other = other._start
+ # check if inside the set
+ if curr_other not in self._map: return False
+ # same start pointer in the set
+ curr = curr_other
+
+ # traverse the subset and check each element appeared in the same order in the set
+ for item in other:
+ while True:
+ # return False if we are at the end
+ if curr == self.NonePtr: return False
+ # if item in the set, go to the next item in the subset
+ if curr == item: break
+ # go to the next item in set
+ curr = self._map[curr][1]
+
+ # return True as we checked that all the items in the subset are in the set
+ return True
+
+ # the order is not important
+ else:
+ return all([(item in self) for item in other])
+
+ # alias
+ contains = issuperset
+
+ def issubset(self, other, order=True):
+ """
+ Return True if the other set is a superset of this set; i.e. return true if this set is a subset of the
+ other set.
+ Time complexity: same as `issuperset()`.
+ """
+ return other.issuperset(self, order=order)
+
+ def __repr__(self):
+ return '%s(%r)' % (self.__class__.__name__, list(self))
+
+ def __contains__(self, item):
+ """
+ Check if the given item/subset is in the set. The subset doesn't have the same order.
+ If the order is important, see the `issubset()` method.
+ Time complexity: O(1) if single item, O(K) if subset (where K is the size of the subset)
+ """
+ if isinstance(item, OrderedSet):
+ return self.issuperset(item, order=False)
+ return item in self._map
+
+ def __len__(self):
+ """
+ Return the length of the ordered set.
+ Time complexity: O(1)
+ """
+ return len(self._map)
+
+ def __iter__(self):
+ """
+ Iterate over the ordered set in the order the items have been added.
+ Time complexity: O(N)
+ """
+ curr = self._start
+ while curr != self.NonePtr:
+ yield curr
+ curr = self._map[curr][1]
+
+ def __reversed__(self):
+ """
+ Iterate over the ordered set in the reverse order the items have been added.
+ Time complexity: O(N)
+ """
+ curr = self._end
+ while curr != self.NonePtr:
+ yield curr
+ curr = self._map[curr][0]
+
+ def __getitem__(self, idx):
+ """
+ Return the item (ordered set) associated to the given index (indices).
+ Time complexity: O(N)
+ """
+ # checks
+ if not isinstance(idx, (int, slice)):
+ raise KeyError("Expecting an int or slice for the index.")
+ if len(self._map) == 0:
+ raise KeyError("Trying to get an item from an empty set.")
+
+ if isinstance(idx, int): # index is an integer
+
+ # check index
+ idx = self._checkIndex(idx)
+
+ # traverse the set in a specific order based on how close the index is wrt the start/end of the set
+ curr = self.NonePtr
+ if 0 <= idx <= len(self._map) / 2: # traverse from the beginning
+ count = 0
+ for curr in self:
+ if idx == count: break
+ count += 1
+ else: # traverse from the end
+ count = len(self._map)-1
+ for curr in reversed(self):
+ if idx == count: break
+ count -= 1
+
+ # return corresponding item
+ return curr
+
+ else: # multiple indices
+ # check arguments from slice
+ lst = []
+ start, stop, step = idx.start, idx.stop, idx.step
+ iterator = self
+ if step is None: step = 1
+ if step > 0:
+ if start is None: start = 0
+ if stop is None: stop = len(self)
+ else:
+ iterator = reversed(self)
+ if start is None: start = len(self) - 1
+ if stop is None: stop = -1
+ start = len(self) - 1 - start
+ stop = len(self) - 1 - stop
+ step = abs(step)
+
+ # traverse the ordered set, and add the requested items into the list
+ count = 0
+ for item in iterator:
+ if count >= stop: break
+ if count < start: pass
+ else:
+ if ((count - start) % step) == 0:
+ lst.append(item)
+ count += 1
+
+ # return a new ordered set
+ return OrderedSet(lst)
+
+ def __setitem__(self, idx, item):
+ """
+ Replace the item at the specified index. If the item is already in the ordered set, it will move it to the
+ specified index.
+ Time complexity: O(N)
+ """
+ # check idx
+ idx = self._checkIndex(idx)
+
+ # check if item already in the set
+ if item in self._map:
+ # move the item at the specified index
+ self.move(idx, item)
+ else: # item is not in the set, thus replace the item at the specified index
+ count = 0
+ for curr in self: # go through the ordered set
+ if count == idx:
+ # add new item
+ prev_item, next_item = self._map[curr]
+ self._map[item] = (prev_item, next_item)
+
+ # set the start/end pointers if needed
+ if count == 0:
+ self._start = item
+ if count == len(self):
+ self._end = item
+
+ # remove the item we have to replaced
+ self._map.pop(curr)
+ break
+ count += 1
+
+ def __add__(self, other):
+ """
+ Union between two ordered sets.
+ """
+ return self | other
+
+ def __iadd__(self, other):
+ """
+ Update a set with the union of itself and the other set.
+ """
+ self |= other
+
+ def __and__(self, other):
+ """
+ Intersection between two ordered sets.
+ """
+ return super(OrderedSet, other).__and__(self)
+
+ def __iand__(self, other):
+ """
+ Update a set with the intersection of itself and the other set.
+ """
+ super(OrderedSet, self).__iand__(other)
+
+ def __mul__(self, other):
+ """
+ Intersection between two ordered sets.
+ """
+ return self & other
+
+ def __imul__(self, other):
+ """
+ Update a set with the intersection of itself and the other set.
+ """
+ self &= other
+
+
+#OrderedSet = OrderedSet2
+
+
+# Tests
+if __name__ == '__main__':
+ # Test the first order set
+ s0 = OrderedSet()
+ l = [9,82,10,-5,-6,4]
+ s1 = OrderedSet(l + [82,10])
+ s2 = OrderedSet([10,9,48,56])
+ s3 = s2 & s1
+
+ print("\nOrdered sets:")
+ print("s0: {}".format(s0))
+ print("s1: {}".format(s1))
+ print("s2: {}".format(s2))
+ print("s3: {}".format(s3))
+
+ print("\nSubsets:")
+ print("82 in s1? {}".format(82 in s1))
+ print("s0 in s0? {}".format(s0 in s0))
+ print("s0 in s1? {}".format(s0 in s1))
+ print("s2 in s1? {}".format(s2 in s1))
+ print("s3 in s1? {}".format(s3 in s1))
+ print("s3 in s2? {}".format(s3 in s2))
+ print("s1.issuperset(s3, order=True) = {}".format(s1.issuperset(s3)))
+ print("s3.issubset(s1) = {}".format(s3.issubset(s1)))
+ print("s1.contains(s3) = {}".format(s1.contains(s3)))
+ print("s2.issuperset(s3, order=True) = {}".format(s2.issuperset(s3)))
+ print("s3.issubset(s2) = {}".format(s3.issubset(s2)))
+ print("s2.contains(s3) = {}".format(s2.contains(s3)))
+
+ print("\nIndexing:")
+ print("s1[0] = {}".format(s1[0]))
+ print("s1[2] = {}".format(s1[2]))
+ print("s1[-1] = {}".format(s1[-1]))
+ print("s1[:4] = {} and l[:4] = {}".format(s1[:4], l[:4]))
+ print("s1[2:4] = {} and l[2:4] = {}".format(s1[2:4], l[2:4]))
+ print("s1[2:] = {} and l[2:] = {}".format(s1[2:], l[2:]))
+ print("s1[::2] = {} and l[::2] = {}".format(s1[::2], l[::2]))
+ print("s1[1:4:2] = {} and l[1:4:2] = {}".format(s1[1:4:2], l[1:4:2]))
+ print("s1[::-1] = {} and l[::-1] = {}".format(s1[::-1], l[::-1]))
+ print("s1[4::-1] = {} and l[4::-1] = {}".format(s1[4::-1], l[4::-1]))
+ print("s1[:1:-1] = {} and l[:1:-1] = {}".format(s1[:1:-1], l[:1:-1]))
+ print("s1[:2:-1] = {} and l[:2:-1] = {}".format(s1[:2:-1], l[:2:-1]))
+ print("s1[4:1:-1] = {} and l[4:1:-1] = {}".format(s1[4:1:-1], l[4:1:-1]))
+ print("s1[4:1:-2] = {} and l[4:1:-2] = {}".format(s1[4:1:-2], l[4:1:-2]))
+
+ print("\nDisjoint:")
+ print("s1.isdisjoint(s0) = {}".format(s1.isdisjoint(s0)))
+ print("s1.isdisjoint(s2) = {}".format(s1.isdisjoint(s2)))
+
+ print("\nUnion:")
+ print("s1 | s2 = {}".format(s1 | s2))
+ print("s1 + s2 = {}".format(s1 + s2))
+ print("s1.union(s2) = {}".format(s1.union(s2)))
+ print("s2 | s1 = {}".format(s2 | s1))
+ print("s2 + s1 = {}".format(s2 + s1))
+ print("s2.union(s1) = {}".format(s2.union(s1)))
+
+ print("\nIntersection:")
+ print("s1 & s2 = {}".format(s1 & s2))
+ print("s1 * s2 = {}".format(s1 * s2))
+ print("s1.intersection(s2) = {}".format(s1.intersection(s2)))
+ print("s2 & s1 = {}".format(s2 & s1))
+ print("s2 * s1 = {}".format(s2 * s1))
+ print("s2.intersection(s1) = {}".format(s2.intersection(s1)))
+
+ print("\nDifference:")
+ print("s1 - s3 = {}".format(s1 - s3))
+ print("s1.difference(s3) = {}".format(s1.difference(s3)))
+ print("s3 - s1 = {}".format(s3 - s1))
+ print("s3.difference(s1) = {}".format(s3.difference(s1)))
+ print("s1 - s2 = {}".format(s1 - s2))
+ print("s1 - s0 = {}".format(s1 - s0))
\ No newline at end of file
diff --git a/pyrobolearn/utils/distributions.py b/pyrobolearn/utils/distributions.py
new file mode 100644
index 0000000..ddbf01d
--- /dev/null
+++ b/pyrobolearn/utils/distributions.py
@@ -0,0 +1,105 @@
+# This file provides the most common probability distributions
+
+# References:
+# [1] pyrobolearn/models/gaussian
+# [2] Distributions in pytorch: https://pytorch.org/docs/stable/distributions.html
+# [3] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr/blob/master/distributions.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+FixedCategorical = torch.distributions.Categorical
+FixedCategorical.sample = lambda self: FixedCategorical.sample(self).unsqueeze(-1)
+FixedCategorical.log_probs = lambda self, actions: FixedCategorical.log_prob(self, actions.squeeze(-1)).unsqueeze(-1)
+FixedCategorical.mode = lambda self: self.probs.argmax(dim=1, keepdim=True)
+
+Normal = torch.distributions.Normal
+Normal.log_probs = lambda self, actions: Normal.log_prob(self, actions).sum(-1, keepdim=True)
+Normal.entropy = lambda self: Normal.entropy(self).sum(-1)
+Normal.mode = lambda self: self.mean
+
+MVN = torch.distributions.MultivariateNormal
+MVN.log_probs = lambda self, actions: MVN.log_prob(self, actions)
+MVN.mode = lambda self: self.mean
+
+
+def init(module, weight_init, bias_init, gain=1):
+ weight_init(module.weight.data, gain=gain)
+ bias_init(module.bias.data)
+ return module
+
+
+# https://github.com/openai/baselines/blob/master/baselines/common/tf_util.py#L87
+def init_normc_(weight, gain=1):
+ # initialize the weights
+ weight.normal_(0, 1)
+ weight *= gain / torch.sqrt(weight.pow(2).sum(1, keepdim=True))
+
+
+class AddBias(nn.Module):
+ def __init__(self, bias):
+ super(AddBias, self).__init__()
+ self._bias = nn.Parameter(bias.unsqueeze(1))
+
+ def forward(self, x):
+ if x.dim() == 2:
+ bias = self._bias.t().view(1, -1)
+ else:
+ bias = self._bias.t().view(1, -1, 1, 1)
+ return x + bias
+
+
+class Categorical(nn.Module):
+ r"""Categorical distribution
+
+ """
+ def __init__(self, num_inputs, num_outputs):
+ super(Categorical, self).__init__()
+
+ init_ = lambda m: init(m, nn.init.orthogonal_, lambda x: nn.init.constant_(x, 0), gain=0.01)
+
+ self.linear = init_(nn.Linear(num_inputs, num_outputs))
+
+ def forward(self, x):
+ x = self.linear(x)
+ return FixedCategorical(logits=x)
+
+
+class DiagonalGaussian(nn.Module):
+ r"""Diagonal Gaussian distribution
+
+ This multivariate gaussian distribution has a diagonal covariance matrix, that is, the variables are independent
+ between each other.
+ """
+ def __init__(self, num_inputs, num_outputs):
+ super(DiagonalGaussian, self).__init__()
+
+ init_ = lambda m: init(m, init_normc_, lambda x: nn.init.constant_(x, 0))
+
+ self.fc_mean = init_(nn.Linear(num_inputs, num_outputs))
+ self.logstd = AddBias(torch.zeros(num_outputs))
+
+ def forward(self, x):
+ action_mean = self.fc_mean(x)
+
+ # An ugly hack for my KFAC implementation.
+ zeros = torch.zeros(action_mean.size())
+ if x.is_cuda:
+ zeros = zeros.cuda()
+
+ action_logstd = self.logstd(zeros)
+ return Normal(action_mean, action_logstd.exp())
+
+
+class FixedDiagonalMVN(nn.Module):
+ r"""Fixed Diagonal Multivariate Normal
+
+ """
+ def __init__(self, num_outputs, variance=1.):
+ super(FixedDiagonalMVN, self).__init__()
+ self.cov = torch.diag(variance * torch.ones(num_outputs))
+
+ def forward(self, x):
+ return MVN(x, covariance_matrix=self.cov)
diff --git a/pyrobolearn/utils/gazebo-ros/spawn_model.py b/pyrobolearn/utils/gazebo-ros/spawn_model.py
new file mode 100755
index 0000000..ccbc337
--- /dev/null
+++ b/pyrobolearn/utils/gazebo-ros/spawn_model.py
@@ -0,0 +1,317 @@
+#!/usr/bin/env python
+#
+# Copyright 2013 Open Source Robotics Foundation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Desc: helper script for spawning models in gazebo
+# Author: John Hsu, Dave Coleman
+#
+
+import rospy, sys, os, time
+import string
+import warnings
+import re
+
+from gazebo_ros import gazebo_interface
+
+from gazebo_msgs.msg import *
+from gazebo_msgs.srv import *
+from std_srvs.srv import Empty
+from geometry_msgs.msg import Point, Pose, Quaternion, Twist, Wrench
+import tf.transformations as tft
+
+model_database_template = """
+
+
+ model://MODEL_NAME
+
+
+"""
+
+def usage():
+ print('''Commands:
+ -[urdf|sdf|trimesh|gazebo] - specify incoming xml is urdf, sdf or trimesh format. gazebo arg is deprecated in ROS Hydro
+ -[file|param|database] [||] - source of the model xml or the trimesh file
+ -model - name of the model to be spawned.
+ -reference_frame - optinal: name of the model/body where initial pose is defined.
+ If left empty or specified as "world", gazebo world frame is used.
+ -gazebo_namespace - optional: ROS namespace of gazebo offered ROS interfaces. Defaults to /gazebo/ (e.g. /gazebo/spawn_model).
+ -robot_namespace - optional: change ROS namespace of gazebo-plugins.
+ -unpause - optional: !!!Experimental!!! unpause physics after spawning model
+ -wait - optional: !!!Experimental!!! wait for model to exist
+ -trimesh_mass - required if -trimesh is used: linear mass
+ -trimesh_ixx - required if -trimesh is used: moment of inertia about x-axis
+ -trimesh_iyy - required if -trimesh is used: moment of inertia about y-axis
+ -trimesh_izz - required if -trimesh is used: moment of inertia about z-axis
+ -trimesh_gravity - required if -trimesh is used: gravity turned on for this trimesh model
+ -trimesh_material - required if -trimesh is used: E.g. Gazebo/Blue
+ -trimesh_name - required if -trimesh is used: name of the link containing the trimesh
+ -x - optional: initial pose, use 0 if left out
+ -y - optional: initial pose, use 0 if left out
+ -z - optional: initial pose, use 0 if left out
+ -R - optional: initial pose, use 0 if left out
+ -P - optional: initial pose, use 0 if left out
+ -Y - optional: initial pose, use 0 if left out
+ -J - optional: initialize the specified joint at the specified value
+ -package_to_model - optional: convert urdf i+2:
+ self.joint_names.append(sys.argv[i+1])
+ self.joint_positions.append(float(sys.argv[i+2]))
+ else:
+ rospy.logerr("Error: must specify a joint name and joint value pair")
+ sys.exit(0)
+ if sys.argv[i] == '-param':
+ if len(sys.argv) > i+1:
+ if self.file_name != "" or self.database_name != "":
+ rospy.logerr("Error: you cannot specify file name if parameter or database name is given, must pick one source of model xml")
+ sys.exit(0)
+ else:
+ self.param_name = sys.argv[i+1]
+ if sys.argv[i] == '-file':
+ if len(sys.argv) > i+1:
+ if self.param_name != "" or self.database_name != "":
+ rospy.logerr("Error: you cannot specify parameter if file or database name is given, must pick one source of model xml")
+ sys.exit(0)
+ else:
+ self.file_name = sys.argv[i+1]
+ if sys.argv[i] == '-database':
+ if len(sys.argv) > i+1:
+ if self.param_name != "" or self.file_name != "":
+ rospy.logerr("Error: you cannot specify parameter if file or parameter name is given, must pick one source of model xml")
+ sys.exit(0)
+ else:
+ self.database_name = sys.argv[i+1]
+ if sys.argv[i] == '-model':
+ if len(sys.argv) > i+1:
+ self.model_name = sys.argv[i+1]
+ if sys.argv[i] == '-wait':
+ if len(sys.argv) > i+1:
+ self.wait_for_model = sys.argv[i+1]
+ if sys.argv[i] == '-reference_frame':
+ if len(sys.argv) > i+1:
+ self.reference_frame = sys.argv[i+1]
+ if sys.argv[i] == '-robot_namespace':
+ if len(sys.argv) > i+1:
+ self.robot_namespace = sys.argv[i+1]
+ if sys.argv[i] == '-namespace':
+ if len(sys.argv) > i+1:
+ self.robot_namespace = sys.argv[i+1]
+ if sys.argv[i] == '-gazebo_namespace':
+ if len(sys.argv) > i+1:
+ self.gazebo_namespace = sys.argv[i+1]
+ if sys.argv[i] == '-x':
+ if len(sys.argv) > i+1:
+ self.initial_xyz[0] = float(sys.argv[i+1])
+ if sys.argv[i] == '-y':
+ if len(sys.argv) > i+1:
+ self.initial_xyz[1] = float(sys.argv[i+1])
+ if sys.argv[i] == '-z':
+ if len(sys.argv) > i+1:
+ self.initial_xyz[2] = float(sys.argv[i+1])
+ if sys.argv[i] == '-R':
+ if len(sys.argv) > i+1:
+ self.initial_rpy[0] = float(sys.argv[i+1])
+ if sys.argv[i] == '-P':
+ if len(sys.argv) > i+1:
+ self.initial_rpy[1] = float(sys.argv[i+1])
+ if sys.argv[i] == '-Y':
+ if len(sys.argv) > i+1:
+ self.initial_rpy[2] = float(sys.argv[i+1])
+ if sys.argv[i] == '-package_to_model':
+ self.package_to_model = True;
+ if sys.argv[i] == '-b':
+ self.bond = True
+
+ if not self.sdf_format and not self.urdf_format:
+ rospy.logerr("Error: you must specify incoming format as either urdf or sdf format xml")
+ sys.exit(0)
+ if self.model_name == "":
+ rospy.logerr("Error: you must specify model name")
+ sys.exit(0)
+
+ def checkForModel(self,model):
+ for n in model.name:
+ if n == self.wait_for_model:
+ self.wait_for_model_exists = True
+
+
+ # Generate a blank SDF file with an include for the model from the model database
+ def createDatabaseCode(self, database_name):
+ return model_database_template.replace("MODEL_NAME", database_name);
+
+ def callSpawnService(self):
+
+ # wait for model to exist
+ rospy.init_node('spawn_model')
+
+ if not self.wait_for_model == "":
+ rospy.Subscriber("%s/model_states"%(self.gazebo_namespace), ModelStates, self.checkForModel)
+ r= rospy.Rate(10)
+ while not rospy.is_shutdown() and not self.wait_for_model_exists:
+ r.sleep()
+
+ if rospy.is_shutdown():
+ sys.exit(0)
+
+ if self.file_name != "":
+ rospy.loginfo("Loading model XML from file")
+ if os.path.exists(self.file_name):
+ if os.path.isdir(self.file_name):
+ rospy.logerr("Error: file name is a path? %s", self.file_name)
+ sys.exit(0)
+ if not os.path.isfile(self.file_name):
+ rospy.logerr("Error: unable to open file %s", self.file_name)
+ sys.exit(0)
+ else:
+ rospy.logerr("Error: file does not exist %s", self.file_name)
+ sys.exit(0)
+ # load file
+ f = open(self.file_name,'r')
+ model_xml = f.read()
+ if model_xml == "":
+ rospy.logerr("Error: file is empty %s", self.file_name)
+ sys.exit(0)
+
+ # ROS Parameter
+ elif self.param_name != "":
+ rospy.loginfo( "Loading model XML from ros parameter")
+ model_xml = rospy.get_param(self.param_name)
+ if model_xml == "":
+ rospy.logerr("Error: param does not exist or is empty")
+ sys.exit(0)
+
+ # Gazebo Model Database
+ elif self.database_name != "":
+ rospy.loginfo( "Loading model XML from Gazebo Model Database")
+ model_xml = self.createDatabaseCode(self.database_name)
+ if model_xml == "":
+ rospy.logerr("Error: an error occured generating the SDF file")
+ sys.exit(0)
+ else:
+ rospy.logerr("Error: user specified param or filename is an empty string")
+ sys.exit(0)
+
+ if self.package_to_model:
+ model_xml = re.sub("<\s*mesh\s+filename\s*=\s*([\"|'])package://","model://", model_xml)
+
+ # setting initial pose
+ initial_pose = Pose()
+ initial_pose.position.x = self.initial_xyz[0]
+ initial_pose.position.y = self.initial_xyz[1]
+ initial_pose.position.z = self.initial_xyz[2]
+ # convert rpy to quaternion for Pose message
+ tmpq = tft.quaternion_from_euler(self.initial_rpy[0],self.initial_rpy[1],self.initial_rpy[2])
+ q = Quaternion(tmpq[0],tmpq[1],tmpq[2],tmpq[3])
+ initial_pose.orientation = q;
+
+ # spawn model
+ if self.urdf_format:
+ success = gazebo_interface.spawn_urdf_model_client(self.model_name, model_xml, self.robot_namespace,
+ initial_pose, self.reference_frame, self.gazebo_namespace)
+ elif self.sdf_format:
+ success = gazebo_interface.spawn_sdf_model_client(self.model_name, model_xml, self.robot_namespace,
+ initial_pose, self.reference_frame, self.gazebo_namespace)
+ else:
+ rospy.logerr("Error: should not be here in spawner helper script, there is a bug")
+ sys.exit(0)
+
+ # set model configuration before unpause if user requested
+ if len(self.joint_names) != 0:
+ try:
+ success = gazebo_interface.set_model_configuration_client(self.model_name, self.param_name,
+ self.joint_names, self.joint_positions, self.gazebo_namespace)
+ except rospy.ServiceException as e:
+ rospy.logerr("Set model configuration service call failed: %s", e)
+
+ # unpause physics if user requested
+ if self.unpause_physics:
+ rospy.wait_for_service('%s/unpause_physics'%(self.gazebo_namespace))
+ try:
+ unpause_physics = rospy.ServiceProxy('%s/unpause_physics'%(self.gazebo_namespace), Empty)
+ unpause_physics()
+ except rospy.ServiceException as e:
+ rospy.logerr("Unpause physics service call failed: %s", e)
+
+ return
+
+
+ def callDeleteService(self):
+ try:
+ delete_model = rospy.ServiceProxy('%s/delete_model'%(self.gazebo_namespace), DeleteModel)
+ delete_model(model_name=self.model_name)
+ except rospy.ServiceException as e:
+ rospy.logerr("Delete model service call failed: %s", e)
+
+if __name__ == "__main__":
+ if len(sys.argv) < 2:
+ print(usage())
+ else:
+ print("SpawnModel script started") # make this a print incase roscore has not been started
+ sm = SpawnModel()
+ sm.parseUserInputs()
+ sm.callSpawnService()
+
+ if sm.bond:
+ rospy.on_shutdown(sm.callDeleteService)
+ rospy.spin()
diff --git a/pyrobolearn/utils/heightmap_generator.py b/pyrobolearn/utils/heightmap_generator.py
new file mode 100644
index 0000000..ff084b3
--- /dev/null
+++ b/pyrobolearn/utils/heightmap_generator.py
@@ -0,0 +1,424 @@
+import numpy as np
+from sklearn.gaussian_process import GaussianProcessRegressor
+from sklearn.gaussian_process.kernels import RBF
+from scipy.interpolate import Rbf
+
+try:
+ import gdal
+except ImportError as e:
+ raise ImportError(repr(e) + '\nTry to install gdal: pip install gdal')
+
+
+def diamond_square_algorithm(N=128, init_values=None, noise=0, lower_bound=0, upper_bound=255, dtype=np.int, seed=None):
+ r"""Diamond-Square Algorithm
+
+ This function implements the diamond-square algorithm [1], to generate random terrains given an initial value
+ for each corner.
+
+ Warnings: the diamond-square algo assumes that the heightmap is a 2D square array.
+
+ Args:
+ N (int): number of points (must be a power of 2). From this, the width and the height will automatically be
+ computed, such that width = height = 2*N+1.
+ init_values (np.array[4], None): the four initial values for the corners. If None, it will generate 4 values
+ randomly such that they are between the lower_bound and upper_bound.
+ noise (int,float): noise level to add. This corresponds to the standard deviation of the normal distribution.
+ lower_bound (int,float): lower bound; each value in the heightmap will be higher than or equal to this bound
+ upper_bound (int,float): upper bound; each value in the heightmap will be lower than or equal to this bound
+ dtype (np.int, np.float): type of the returned array for the heightmap
+ seed (int, None): random seed
+
+ Returns:
+ np.array[2*N+1,2*N+1]: resulting 2D square heightmap
+
+ References:
+ [1] Wikipedia: https://en.wikipedia.org/wiki/Diamond-square_algorithm
+ [2] https://blog.habrador.com/2013/02/how-to-generate-random-terrain.html
+ """
+ # set the seed if given
+ if seed:
+ np.random.seed(seed)
+
+ # create initial heightmap
+ width, height = 2 * N + 1, 2 * N + 1
+ heightmap = -1 * np.ones((height, width), dtype=dtype)
+ if not init_values:
+ if dtype == np.int:
+ init_values = np.random.randint(low=lower_bound, high=upper_bound+1, size=4)
+ else:
+ init_values = np.random.uniform(low=lower_bound, high=upper_bound, size=4)
+ heightmap[0, 0], heightmap[0, width - 1], heightmap[height - 1, 0], heightmap[height - 1, width - 1] = init_values
+
+ # define diamond-square step function
+ def diamond_square_step(heightmap, square=None, noise=0, lower_bound=0, upper_bound=255):
+ """
+ Diamond-square step which which performs a diamond step followed by a square step.
+
+ Args:
+ heightmap (np.array[2*N+1,2*N+1]): heightmap (initial square)
+ square (np.array[M,M]): the current square we focus on.
+ """
+ # if no square given
+ if square is None:
+ height, width = heightmap.shape
+ square = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]])
+
+ # check size of square
+ xmin, xmax, ymin, ymax = square[:, 0].min(), square[:, 0].max(), square[:, 1].min(), square[:, 1].max()
+ dx, dy = (xmax - xmin), (ymax - ymin)
+ if dx == 0 or dx == 1 or dy == 0 or dy == 1:
+ return
+
+ # DIAMOND STEP
+ center = np.array([xmin + dx / 2, ymin + dy / 2])
+ yc, xc = center
+ heightmap[xc, yc] = np.mean([heightmap[x, y] for (y, x) in square]) #+ np.random.normal(scale=noise)
+ heightmap[xc, yc] = min(max(lower_bound, heightmap[xc, yc]), upper_bound) # lower and upper bound
+
+ # SQUARE STEP
+ # triangles: a triangle is defined by 3 points
+ triangles = np.array([[c1, c2, center] for c1, c2 in zip(square, list(square[1:]) + [square[0]])])
+
+ squares = []
+ for i, triangle in enumerate(triangles):
+ xmin, xmax, ymin, ymax = triangle[:, 0].min(), triangle[:, 0].max(), triangle[:, 1].min(), triangle[:,
+ 1].max()
+
+ if i == 0: # upper triangle
+ center = np.array([xmin + (xmax - xmin) / 2, ymin])
+ square = np.array([[xmin, ymin], center, [center[0], ymax], [xmin, ymax]]) # left upper square
+ elif i == 1: # right triangle
+ center = np.array([xmax, ymin + (ymax - ymin) / 2])
+ square = np.array([[xmin, ymin], [xmax, ymin], center, [xmin, center[1]]]) # right upper square
+ elif i == 2: # lower triangle
+ center = np.array([xmin + (xmax - xmin) / 2, ymax])
+ square = np.array([[center[0], ymin], [xmax, ymin], [xmax, ymax], center]) # right lower square
+ else: # left triangle
+ center = np.array([xmin, ymin + (ymax - ymin) / 2])
+ square = np.array([center, [xmax, center[1]], [xmax, ymax], [xmin, ymax]]) # left lower square
+
+ yc, xc = center
+ heightmap[xc, yc] = np.mean([heightmap[x, y] for (y, x) in triangle]) #+ np.random.normal(scale=noise)
+ heightmap[xc, yc] = min(max(lower_bound, heightmap[xc, yc]), upper_bound) # lower and upper bound
+
+ # a square is defined by 4 points
+ squares.append(square)
+
+ # for each subsquare in the original square, compute the heightmap recursively
+ for square in squares:
+ diamond_square_step(heightmap, square, noise, lower_bound, upper_bound)
+
+ # start diamond-square algorithm (recursively)
+ diamond_square_step(heightmap, noise=noise, lower_bound=lower_bound, upper_bound=upper_bound)
+ return heightmap
+
+
+
+def heightmap_gpr(init_values, x, y, kernel=None, alpha=1e-10, lower_bound=0, upper_bound=255, dtype=np.int):
+ r"""
+ Generate a heightmap using gaussian process regression. The advantages of using this method over others to
+ generate terrains lies in the capacity of adding prior knowledge through the kernel and the given initial values.
+ For instance, using a RBF kernel means that we want a smooth terrain instead of a bumpy one.
+ Furthermore, it allows to generate heightmaps which are not necessary square; i.e. they can be rectangular.
+
+ Warnings: this is pretty difficult to exploit if the given data is not consistent. See `heigthmap_rbf` for
+ a better way to generate heightmap.
+
+ Args:
+ init_values (np.array[M,3]): list of `M` 3D points which corresponds to initial values that are used to fit
+ the gaussian process.
+ x (np.array[N], np.array[N,O]): If 1d array, it will compute the meshgrid. Otherwise, the resulting 2D array
+ from the meshgrid is expected. This is used to predict the heightmap at the given points.
+ y (np.array[0], np.array[N,O]): If 1d array, it will compute the meshgrid. Otherwise, the resulting 2D array
+ from the meshgrid is expected. This is used to predict the heightmap at the given points.
+ kernel (None, sklearn.gaussian_process.kernels.Kernel): "The kernel specifying the covariance function of
+ the GP. If None is passed, the kernel '1.0 * RBF(1.0)' is used as default. Note that the kernel's
+ hyperparameters are optimized during fitting" [2]
+ alpha (float, array_like): "Value added to the diagonal of the kernel matrix during fitting. Larger values
+ correspond to increased noise level in the observations. This can also prevent a potential numerical issue
+ during fitting, by ensuring that the calculated values form a positive definite matrix. If an array is
+ passed, it must have the same number of entries as the data used for fitting and is used as
+ datapoint-dependent noise level. Note that this is equivalent to adding a WhiteKernel with c=alpha.
+ Allowing to specify the noise level directly as a parameter is mainly for convenience and for consistency
+ with Ridge." [2]
+ lower_bound (int,float): lower bound; each value in the heightmap will be higher than or equal to this bound
+ upper_bound (int,float): upper bound; each value in the heightmap will be lower than or equal to this bound
+ dtype (np.int, np.float): type of the returned array for the heightmap
+
+ Returns:
+ np.array[N,O]: resulting 2D heightmap
+
+ References:
+ [1] "Gaussian Processes for Machine Learning", Rasmussen and Williams, 2006
+ [2] Sklearn: https://scikit-learn.org/stable/modules/gaussian_process.html
+ """
+ # check given x and y
+ if len(x.shape) == 1 and len(y.shape) == 1:
+ x,y = np.meshgrid(x,y)
+ if x.shape != y.shape:
+ raise ValueError("Expecting x and y to have the same shape, which should be the case if it is a meshgrid")
+
+ # compute the minimum distance between points
+ N = len(init_values)
+ min_dist = np.inf
+ for i in range(N):
+ for j in range(i+1,N):
+ dist = np.linalg.norm(init_values[i,:2] - init_values[j,:2])
+ if dist < min_dist:
+ min_dist = dist
+ print("Min dist: {}".format(min_dist))
+
+ # check initial values
+ if not isinstance(init_values, np.ndarray):
+ raise TypeError("Expecting init_values to be a numpy array")
+ if init_values.shape[1] != 3:
+ raise ValueError("Expecting a numpy array of 3D points for init_values")
+
+ # create gaussian process and fit on the given initial values
+ kernel = RBF(length_scale=np.sqrt(min_dist))
+ gpr = GaussianProcessRegressor(kernel=kernel, alpha=alpha, normalize_y=True)
+ gpr.fit(init_values[:,:2], init_values[:,2])
+
+ # predict the heightmap using GPR
+ X = np.dstack((x,y)).reshape(-1,2)
+ heightmap = gpr.predict(X)
+ heightmap = heightmap.reshape(x.shape)
+
+ print("Params: {}".format(gpr.get_params()))
+
+ # make sure the values of the heightmap are between the bounds (in-place), and is the correct type
+ np.clip(heightmap, lower_bound, upper_bound, heightmap)
+ heightmap.astype(dtype)
+
+ return heightmap
+
+
+
+def heightmap_rbf(init_values, x, y, function='multiquadric', lower_bound=0, upper_bound=255, dtype=np.int):
+ r"""
+ Generate heightmap by interpolating the given initial points using RBF functions.
+
+ Advantages: fast and easy to use, and the results are pretty good. Heightmaps can also be rectangular.
+
+ Args:
+ init_values (np.array[M,3]): list of `M` 3D points which corresponds to initial values that are used to fit
+ the gaussian process.
+ x (np.array[N], np.array[N,O]): If 1d array, it will compute the meshgrid. Otherwise, the resulting 2D array
+ from the meshgrid is expected. This is used to predict the heightmap at the given points.
+ y (np.array[0], np.array[N,O]): If 1d array, it will compute the meshgrid. Otherwise, the resulting 2D array
+ from the meshgrid is expected. This is used to predict the heightmap at the given points.
+ function (str, callable): "The radial basis function, based on the radius, r, given by the norm
+ (default is Euclidean distance);
+ 'multiquadric': sqrt((r/self.epsilon)**2 + 1)
+ 'inverse': 1.0/sqrt((r/self.epsilon)**2 + 1)
+ 'gaussian': exp(-(r/self.epsilon)**2)
+ 'linear': r
+ 'cubic': r**3
+ 'quintic': r**5
+ 'thin_plate': r**2 * log(r)
+ If callable, then it must take 2 arguments (self, r). The epsilon parameter will be available as
+ self.epsilon. Other keyword arguments passed in will be available as well." [1]
+ lower_bound (int,float): lower bound; each value in the heightmap will be higher than or equal to this bound
+ upper_bound (int,float): upper bound; each value in the heightmap will be lower than or equal to this bound
+ dtype (np.int, np.float): type of the returned array for the heightmap
+
+ Returns:
+ np.array[N,O]: resulting 2D heightmap
+
+ References:
+ [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.Rbf.html
+ """
+ # check given x and y
+ if len(x.shape) == 1 and len(y.shape) == 1:
+ x, y = np.meshgrid(x, y)
+ if x.shape != y.shape:
+ raise ValueError("Expecting x and y to have the same shape, which should be the case if it is a meshgrid")
+ origin_shape = x.shape
+
+ rbf = Rbf(init_values[:,0], init_values[:,1], init_values[:,2], function=function)
+ heightmap = rbf(x.reshape(-1), y.reshape(-1))
+ heightmap = heightmap.reshape(origin_shape)
+
+ # make sure the values of the heightmap are between the bounds (in-place), and is the correct type
+ np.clip(heightmap, lower_bound, upper_bound, heightmap)
+ heightmap.astype(dtype)
+
+ return heightmap
+
+
+def heighmap_equation(x, y, z, lower_bound=0, upper_bound=255, dtype=np.int):
+ r"""
+ Generate heightmap from 3D equation :math:`z = f(x,y)`.
+
+ Args:
+ x (np.array[N], np.array[N,O]): If 1d array, it will compute the meshgrid. Otherwise, the resulting 2D array
+ from the meshgrid is expected. This is used to predict the heightmap at the given points.
+ y (np.array[0], np.array[N,O]): If 1d array, it will compute the meshgrid. Otherwise, the resulting 2D array
+ from the meshgrid is expected. This is used to predict the heightmap at the given points.
+ z (callable): it must be a function that accepts two arguments `x` and `y` which will be the arrays from the
+ meshgrid.
+ lower_bound (int,float): lower bound; each value in the heightmap will be higher than or equal to this bound
+ upper_bound (int,float): upper bound; each value in the heightmap will be lower than or equal to this bound
+ dtype (np.int, np.float): type of the returned array for the heightmap
+
+ Examples of 2D surfaces:
+ z = lambda x,y: np.log(y)
+ z = lambda x,y: np.sin(np.pi * x) * np.sin(np.pi * y)
+
+ Returns:
+ np.array[N,O]: resulting 2D heightmap
+ """
+ # check given x and y
+ if len(x.shape) == 1 and len(y.shape) == 1:
+ x, y = np.meshgrid(x, y)
+ if x.shape != y.shape:
+ raise ValueError("Expecting x and y to have the same shape, which should be the case if it is a meshgrid")
+ origin_shape = x.shape
+
+ # call z function: z=f(x,y)
+ heightmap = z(x,y)
+
+ # make sure the values of the heightmap are between the bounds (in-place), and is the correct type
+ np.clip(heightmap, lower_bound, upper_bound, heightmap)
+ heightmap.astype(dtype)
+
+ return heightmap
+
+
+def heightmap_gdal(filename, subsample=None, interpolate_fct='multiquadric', lower_bound=0, upper_bound=255,
+ dtype=np.int):
+ r"""
+ Heightmap generated using the Geospatial Data Abstraction Library (GDAL), which allows to open Digital Elevation
+ Models (DEM) or Geographic Information System (GIS). It can open a .tiff, .geotiff, ascii grid, or
+ image (jpg, png,...) file.
+
+ Args:
+ filename (str): path to a DEM, GIS, or image file
+ subsample (int, None): if not None, it is the number of points to sub-sample (to smooth the heightmap using
+ the specified function)
+ interpolate_fct (str, callable): "The radial basis function, based on the radius, r, given by the norm
+ (default is Euclidean distance);
+ 'multiquadric': sqrt((r/self.epsilon)**2 + 1)
+ 'inverse': 1.0/sqrt((r/self.epsilon)**2 + 1)
+ 'gaussian': exp(-(r/self.epsilon)**2)
+ 'linear': r
+ 'cubic': r**3
+ 'quintic': r**5
+ 'thin_plate': r**2 * log(r)
+ If callable, then it must take 2 arguments (self, r). The epsilon parameter will be available as
+ self.epsilon. Other keyword arguments passed in will be available as well." [1]
+ lower_bound (int,float): lower bound; each value in the heightmap will be higher than or equal to this bound
+ upper_bound (int,float): upper bound; each value in the heightmap will be lower than or equal to this bound
+ dtype (np.int, np.float): type of the returned array for the heightmap
+
+ Returns:
+ np.array[H,W]: resulting 2D array of size width `W` and height `H`
+
+ References:
+ [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.Rbf.html
+ """
+ # load data (raster)
+ data = gdal.Open(filename)
+ band = data.GetRasterBand(1)
+ heightmap = band.ReadAsArray() # elevation values
+
+ if isinstance(subsample, int) and subsample > 0:
+ height, width = heightmap.shape
+ idx_x = np.linspace(0, height-1, subsample, dtype=np.int)
+ idx_y = np.linspace(0, width-1, subsample, dtype=np.int)
+ idx_x, idx_y = np.meshgrid(idx_x, idx_y)
+ x,y = np.arange(width), np.arange(height)
+ x,y = np.meshgrid(x,y)
+ rbf = Rbf(x[idx_x, idx_y], y[idx_x, idx_y], heightmap[idx_x, idx_y], function=interpolate_fct)
+ #Nx, Ny = x.shape[0] / subsample, x.shape[1] / subsample
+ #rbf = Rbf(x[::Nx, ::Ny], y[::Nx, ::Ny], heightmap[::Nx, ::Ny], function=interpolate_fct)
+ heightmap = rbf(x, y)
+
+ # make sure the values of the heightmap are between the bounds (in-place), and is the correct type
+ if lower_bound and upper_bound:
+ np.clip(heightmap, lower_bound, upper_bound, heightmap)
+ elif lower_bound:
+ np.clip(heightmap, lower_bound, heightmap.max(), heightmap)
+ elif upper_bound:
+ np.clip(heightmap, heightmap.min(), upper_bound, heightmap)
+ if dtype:
+ heightmap.astype(dtype)
+
+ return heightmap
+
+
+# alias
+heigtmap_from_image = heightmap_gdal
+
+
+
+# Tests
+# Conclusion: use `heightmap_rbf` or `heightmap_gdal` as it is pretty good
+if __name__ == '__main__':
+ from mpl_toolkits.mplot3d import Axes3D
+ import matplotlib.pyplot as plt
+
+ # define plot figure for heightmap
+ def plot_figure(heightmap, title='', block=True, z_upper_lim=256):
+ fig = plt.figure()
+ fig.suptitle(title)
+
+ # 1st subplot: 2D heightmap
+ ax = fig.add_subplot(1, 2, 1)
+ ax.set_title('2D heightmap')
+ ax.imshow(heightmap, cmap='gray')
+
+ # 2nd subplot: associated 3D terrain
+ ax = fig.add_subplot(1, 2, 2, projection='3d')
+ ax.set_title('3D terrain')
+ x = np.linspace(0, 1, heightmap.shape[0])
+ y = np.linspace(0, 1, heightmap.shape[1])
+ x, y = np.meshgrid(y, x)
+ ax.plot_surface(x, y, heightmap)
+ ax.set_zlim(0, z_upper_lim)
+ print(x.shape)
+
+ plt.show(block=block)
+
+
+ # # generate heightmap using the diamond-square algorithm
+ # N = 128 # shape of map: 2N+1, 2N+1
+ # heightmap = diamond_square_algorithm(N)
+ # plot_figure(heightmap, title='Diamond-Square algorithm')
+
+
+ # # generate heightmap using gaussian process regression
+ # x = np.array(range(256))
+ # y = np.array(range(256))
+ # N_init = 20
+ # x_init = np.random.randint(low=x.min(), high=x.max(), size=N_init)
+ # y_init = np.random.randint(low=y.min(), high=y.max(), size=N_init)
+ # z_init = np.random.randint(low=0, high=20, size=N_init)
+ # init_values = np.vstack((x_init, y_init, z_init)).T # shape: Nx3
+ # #init_values = np.array([[163, 73, 0], [13, 15, 1],[69, 102, 2]])
+ # #init_values = np.array([[182, 48, 89], [182, 20, 150], [167, 247, 131]])
+ # heightmap = heightmap_gpr(init_values=init_values, x=x, y=y)
+ # plot_figure(heightmap, title='Gaussian Process Regression')
+
+
+ # generate heightmap using RBF interpolations
+ x = np.array(range(256))
+ y = np.array(range(256)) # range(128)
+ N_init = 20 # number of bumps
+ x_init = np.random.randint(low=x.min(), high=x.max(), size=N_init)
+ y_init = np.random.randint(low=y.min(), high=y.max(), size=N_init)
+ z_init = np.random.randint(low=0, high=20, size=N_init)
+ init_values = np.vstack((x_init, y_init, z_init)).T # shape: Nx3
+ # init_values = np.array([[211, 184, 3], [97, 59, 4], [37, 179, 8], [168, 32, 8], [198, 74, 13],
+ # [44, 10, 2], [175, 102, 6], [6, 22, 1], [35, 165, 6], [169, 211, 16],
+ # [158, 119, 18], [228, 63, 13], [40, 62, 15], [76, 221, 10], [1, 113, 10],
+ # [178, 194, 2], [23, 176,10], [231, 88, 7], [247, 209, 6], [72, 94, 2]])
+ heightmap = heightmap_rbf(init_values=init_values, x=x, y=y, function='gaussian') # 'linear', 'multiquadric'
+ plot_figure(heightmap, title='RBF interpolation')
+
+
+ # generate heigthmap from an image or tif file
+ #dem = heightmap_gdal('../tests/canyon-geo.tif')
+ #dem = heightmap_gdal('../tests/dem.jpg')
+ dem = heightmap_gdal('../tests/heightmap.png')
+ plot_figure(dem, block=True)
\ No newline at end of file
diff --git a/pyrobolearn/utils/human_kinematic.py b/pyrobolearn/utils/human_kinematic.py
new file mode 100644
index 0000000..a3e8787
--- /dev/null
+++ b/pyrobolearn/utils/human_kinematic.py
@@ -0,0 +1,9 @@
+# This file contains a description of the kinematics of a human being
+# You can get link positions/orientations wrt to the world or any links,
+# joint positions (=link orientation wrt previous link), marker positions if any,
+#
+
+class HumanKinematicSkeleton(object):
+
+ def __init__(self):
+ pass
\ No newline at end of file
diff --git a/pyrobolearn/utils/interpolator.py b/pyrobolearn/utils/interpolator.py
new file mode 100644
index 0000000..ea47d8f
--- /dev/null
+++ b/pyrobolearn/utils/interpolator.py
@@ -0,0 +1,110 @@
+
+import numpy as np
+
+
+class HermiteInterpolator(object):
+ r"""5th order Hermite interpolator
+
+ """
+
+ def __init__(self, t, x):
+ """Calculate the coefficients for the interpolation.
+
+ Assuming a trajectory x(t) is described by a fifth order polynomial such that:
+ .. math:: x(t) = a_5 t^5 + a_4 t^4 + a_3 t^3 + a_2 t^2 + a_1 t + a_0
+
+ then taking the derivatives with respect to time give us:
+ .. math::
+ \dot{x}(t) = 5 a_5 t^4 + 4 a_4 t^3 + 3 a_3 t^2 + 2 a_2 t + a_1
+ \ddot{x}(t) = 20 a_5 t^3 + 12 a_4 t^2 + 6 a_3 t + 2 a_2
+
+ We further impose that the initial/final velocities/accelerations to be equal to 0, that is
+ :math:`\dot{x}(t_0) = 0, \dot{x}(t_f) = 0, \ddot{x}(t_0) = 0, \ddot{x}(t_f) = 0`.
+
+ Args:
+ t (float[T]): time
+ x (float[T]): signal/trajectory x(t) to interpolate
+ """
+ if not isinstance(t, (np.ndarray, list, tuple)):
+ raise TypeError("Expecting an iterable for variable t")
+ if not isinstance(x, (np.ndarray, list, tuple)):
+ raise TypeError("Expecting an iterable for variable x")
+
+ tf = t[-1]
+ A = np.array([[1, 1, 1, 1, 1, 1],
+ [5, 4, 3, 2, 1, 0],
+ [20, 12, 6, 2, 0, 0],
+ [0, 0, 0, 1, 0, 0],
+ [0, 0, 0, 0, 1, 0],
+ [0, 0, 0, 0, 0, 1]], dtype=np.float64)
+ A *= np.array([tf**i for i in range(5,-1,-1)])
+ L = len(t) - 2
+ if L != 0:
+ l = []
+ for i in t[1:-1]:
+ l.append([i**j for j in range(5,-1,-1)])
+ A = np.vstack((A, np.array(l)))
+ b = np.array([x[-1], 0, 0, 0, 0, x[0]] + list(x[1:-1]))
+ else:
+ b = np.array([x[-1], 0, 0, 0, 0, x[0]])
+ #coeff = np.linalg.solve(A,b)[0]
+ self.coeff = np.linalg.lstsq(A, b, rcond=None)[0]
+
+ def __call__(self, t):
+ """Interpolate the function.
+
+ Args:
+ t (float, float[T]): time
+
+ Returns:
+ float, float[T]: position
+ float, float[T]: velocity
+ float, float[T]: acceleration
+ """
+ x = np.sum(self.coeff * np.array([[ti**i for i in range(5,-1,-1)] for ti in t]), axis=1)
+ xd = np.sum(self.coeff[:-1] * np.array([[5*ti**4, 4*ti**3, 3*ti**2, 2*ti, 1] for ti in t]), axis=1)
+ xdd = np.sum(self.coeff[:-2] * np.array([[20*ti**3, 12*ti**2, 6*ti, 2] for ti in t]), axis=1)
+ return x, xd, xdd
+
+
+if __name__ == '__main__':
+ import matplotlib.pyplot as plt
+ import matplotlib.gridspec as gridspec
+
+ # define few points in the x-y plane parametrized by t
+ t = np.array([0.0, 0.25, 0.5, 0.75, 1.0])
+ x = np.array([0.5, 0.25, 0.5, 0.75, 0.5])
+ y = np.array([1.0, 0.75, 0.5, 0.25, 0.0])
+
+ # create 5th order Hermite interpolators
+ x_interpolator = HermiteInterpolator(t, x)
+ y_interpolator = HermiteInterpolator(t, y)
+
+ # interpolate the data
+ t = np.linspace(0., 1., 100)
+ x,xd,xdd = x_interpolator(t)
+ y,yd,ydd = y_interpolator(t)
+
+ # plot figures
+ gs = gridspec.GridSpec(4,4)
+ plt.subplot(gs[0, 1:3])
+ plt.title('Hermite Interpolator')
+ plt.plot(x,y)
+ plt.xlabel('x(t)')
+ plt.ylabel('y(t)')
+
+ y_labels = ['x(t)', 'y(t)', 'dx/dt', 'dy/dt', 'd^2x/dt^2', 'd^2y/dt^2']
+ for i, (x_traj, y_traj) in enumerate(zip([x, xd, xdd], [y, yd, ydd])):
+ plt.subplot(gs[i+1, :2])
+ plt.plot(t, x_traj)
+ plt.ylabel(y_labels[2*i])
+ if i == 2:
+ plt.xlabel('t')
+ plt.subplot(gs[i+1, 2:])
+ plt.plot(t, y_traj)
+ plt.ylabel(y_labels[2*i+1])
+ if i == 2:
+ plt.xlabel('t')
+
+ plt.tight_layout()
+ plt.show()
\ No newline at end of file
diff --git a/pyrobolearn/utils/math_utils.py b/pyrobolearn/utils/math_utils.py
new file mode 100644
index 0000000..b3b8c3f
--- /dev/null
+++ b/pyrobolearn/utils/math_utils.py
@@ -0,0 +1,90 @@
+# This file defines mathematical operations
+
+import numpy as np
+import copy
+
+def exp(x):
+ if callable(x):
+ y = copy.copy(x)
+ def exp():
+ return np.exp(x())
+ y.__call__ = exp
+ return y
+ else:
+ return np.exp(x)
+
+
+class Plane(object):
+ """Plane class.
+
+ A plane is defined by its initial point and its normal vector.
+ .. math:: \pi \equiv \overline{n} \cdot (\overline{x} - \overline{x}_0) = 0
+ where :math:`\cdot` is the scalar product operator, :math:`\overline{n}` is the normal vector to the plane
+ :math:`\pi`, :math:`\overline{x_0}` is the initial point on the plane, and :math:`\overline{x}` is an arbitrary
+ point on the plane. Basically, this equation states that any vector on the plane is perpendicular to the normal
+ vector.
+
+ Given a 3D point in the space :math:`\overline{x}_1 = [x_1,y_1,z_1]`, if you wish to know the intersection of
+ the line perpendicular to the plane :math:`\pi` and passing through this point, you can use the fact that this
+ intersection point :math:`\overline{x} = [x,y,z]` has to satisfy the line and plane equations.
+ That is, the line is given by :math:`\overline{x} &= \overline{x}_1 + \lambda \overline{n}`, and by replacing
+ it in the plane equation, and solving it for :math:`\lambda`, and then finally re-incorporating this one into
+ the line equation will give you:
+ .. math:: `\overline{x} = \overline{x}_1 + \overline{n} \cdot (\overline{x}_0 - \overline{x}_1) \overline{n}`
+ """
+
+ def __init__(self, x0, normal):
+ self.threshold = 1e-12
+ self.x0 = x0
+ self.normal = normal
+
+ def convertToArray(self, pt):
+ if isinstance(pt, (tuple, list)):
+ pt = np.array(pt)
+ if not isinstance(pt, np.ndarray):
+ raise TypeError("Expecting a numpy array of shape 3")
+ else:
+ if len(pt.shape) > 1:
+ raise ValueError("Expecting an array")
+ if pt.shape != (3,):
+ raise ValueError("Expecting a shape 3")
+ return pt
+
+ @property
+ def x0(self):
+ return self._x0
+
+ @x0.setter
+ def x0(self, x0):
+ self._x0 = self.convertToArray(x0)
+
+ @property
+ def normal(self):
+ return self._normal
+
+ @normal.setter
+ def normal(self, normal):
+ normal = self.convertToArray(normal)
+ # normalize
+ norm = np.linalg.norm(normal)
+ if norm < self.threshold:
+ raise ValueError("The norm of the normal vector is too close to zero.")
+ self._normal = normal / norm
+
+ def __contains__(self, point):
+ """Check if the given point is in the plane."""
+ point = self.convertToArray(point)
+
+ # scalar product between the normal and (point-x0) vectors
+ val = self.normal.T.dot(point - self.x0)
+
+ if val < self.threshold:
+ return True
+ return False
+
+ def getIntersectionPoint(self, point):
+ """
+ Get the intersection of the plane with a line that starts at the given point and is parallel to the normal.
+ """
+ point = self.convertToArray(point)
+ return point + self.normal.T.dot(self.x0 - point) * self.normal
\ No newline at end of file
diff --git a/pyrobolearn/utils/mesh.py b/pyrobolearn/utils/mesh.py
new file mode 100644
index 0000000..564e62f
--- /dev/null
+++ b/pyrobolearn/utils/mesh.py
@@ -0,0 +1,714 @@
+
+import numpy as np
+try:
+ from mayavi import mlab
+except ImportError as e:
+ raise ImportError(repr(e) + '\nTry to install Mayavi: pip install mayavi')
+try:
+ import gdal
+except ImportError as e:
+ raise ImportError(repr(e) + '\nTry to install gdal: pip install gdal')
+
+import subprocess
+import fileinput
+import sys
+import os
+import scipy.interpolate
+
+
+def recenter(coords):
+ """
+ Recenter the data.
+
+ Args:
+ coords (list of np.array[N], np.array[N]): coordinate(s) to recenter
+
+ Returns:
+ list of np.array[N], np.array[N]: recentered coordinate(s)
+ """
+ if isinstance(coords, (list, tuple)) or len(coords.shape) > 1:
+ centered_coords = []
+ for coord in coords:
+ c_min, c_max = coord.min(), coord.max()
+ c_center = c_min + (c_max - c_min) / 2.
+ centered_coord = coord - c_center
+ centered_coords.append(centered_coord)
+ return np.array(centered_coords)
+
+ c_min, c_max = coords.min(), coords.max()
+ c_center = c_min + (c_max - c_min) / 2.
+ return (coords - c_center)
+
+
+def createMesh(x, y, z, filename=None, show=False, center=True):
+ """
+ Create mesh from x,y,z arrays, and save it in the obj format.
+
+ Args:
+ x (float[N,M]): 2D array representing the x coordinates for the mesh
+ y (float[N,M]): 2D array representing the y coordinates for the mesh
+ z (float[N,M]): 2D array representing the z coordinates for the mesh
+ filename (str, None): filename to save the mesh. If None, it won't save it.
+ show (bool): if True, it will show the mesh using `mayavi.mlab`.
+ center (bool): if True, it will center the mesh
+
+ Examples:
+ # create ellipsoid
+ import numpy as np
+
+ a,b,c,n = 2., 1., 1., 100
+ theta, phi = np.meshgrid(np.linspace(-np.pi/2, np.pi/2, n), np.linspace(-np.pi, np.pi, n))
+
+ x, y, z = a * np.cos(theta) * np.cos(phi), b * np.cos(theta) * np.sin(phi), c * np.sin(theta)
+
+ createMesh(x, y, z, show=True)
+ """
+ #if not (isinstance(x, np.ndarray) and isinstance(y, np.ndarray) and isinstance(z, np.ndarray)):
+ # raise TypeError("Expecting x, y, and z to be numpy arrays")
+
+ if isinstance(x, list) and isinstance(y, list) and isinstance(z, list):
+ # create several 3D mesh
+ for i,j,k in zip(x,y,z):
+ # if we need to recenter
+ if center:
+ i,j,k = recenter([i,j,k])
+ mlab.mesh(i,j,k)
+ else:
+ # if we need to recenter the data
+ if center:
+ x,y,z = recenter([x,y,z])
+
+ # create 3D mesh
+ mlab.mesh(x,y,z)
+
+ # save mesh
+ if filename is not None:
+ if filename[-4:] == '.obj': # This is because the .obj saved by Mayavi is not correct (see in Meshlab)
+ x3dfile = filename[:-4] + '.x3d'
+ mlab.savefig(x3dfile)
+ convertX3dToObj(x3dfile, removeX3d=True)
+ else:
+ mlab.savefig(filename)
+
+ # show / close
+ if show:
+ mlab.show()
+ else:
+ mlab.close()
+
+
+def createSurfMesh(surface, filename=None, show=False, subsample=None, interpolate_fct='multiquadric',
+ lower_bound=None, upper_bound=None, dtype=None):
+ """
+ Create surface (heightmap) mesh, and save it in the obj format.
+
+ Args:
+ surface (float[M,N], str): 2D array where each value represents the height. If it is a string, it is assumed
+ that is the path to a file .tif, .geotiff or an image (.png, .jpg, etc). It will be opened using the
+ `gdal` library.
+ filename (str, None): filename to save the mesh. If None, it won't save it.
+ show (bool): if True, it will show the mesh using `mayavi.mlab`.
+ subsample (int, None): if not None, it is the number of points to sub-sample (to smooth the heightmap using
+ the specified function)
+ interpolate_fct (str, callable): "The radial basis function, based on the radius, r, given by the norm
+ (default is Euclidean distance);
+ 'multiquadric': sqrt((r/self.epsilon)**2 + 1)
+ 'inverse': 1.0/sqrt((r/self.epsilon)**2 + 1)
+ 'gaussian': exp(-(r/self.epsilon)**2)
+ 'linear': r
+ 'cubic': r**3
+ 'quintic': r**5
+ 'thin_plate': r**2 * log(r)
+ If callable, then it must take 2 arguments (self, r). The epsilon parameter will be available as
+ self.epsilon. Other keyword arguments passed in will be available as well." [1]
+ lower_bound (int, float, None): lower bound; each value in the heightmap will be higher than or equal to
+ this bound
+ upper_bound (int, float, None): upper bound; each value in the heightmap will be lower than or equal to
+ this bound
+ dtype (np.int, np.float, None): type of the returned array for the heightmap
+
+ Examples:
+ # create heightmap
+ import numpy as np
+
+ height = np.random.rand(100,100) # in meters
+ createSurfMesh(height, show=True)
+ """
+ if isinstance(surface, str):
+ from utils.heightmap_generator import heightmap_gdal
+ surface = heightmap_gdal(surface, subsample=subsample, interpolate_fct=interpolate_fct,
+ lower_bound=lower_bound, upper_bound=upper_bound, dtype=dtype)
+
+ if not isinstance(surface, np.ndarray):
+ raise TypeError("Expecting a 2D numpy array")
+ if len(surface.shape) != 2:
+ raise ValueError("Expecting a 2D numpy array")
+
+ # create surface mesh
+ mlab.surf(surface)
+
+ # save mesh
+ if filename is not None:
+ if filename[-4:] == '.obj': # This is because the .obj saved by Mayavi is not correct (see in Meshlab)
+ x3dfile = filename[:-4] + '.x3d'
+ mlab.savefig(x3dfile)
+ convertX3dToObj(x3dfile, removeX3d=True)
+ else:
+ mlab.savefig(filename)
+
+ # show / close
+ if show:
+ mlab.show()
+ else:
+ mlab.close()
+
+
+def create3DMesh(heightmap, x=None, y=None, depth_level=1., filename=None, show=False, subsample=None,
+ interpolate_fct='multiquadric', lower_bound=None, upper_bound=None, dtype=None, center=True):
+ """
+ Create 3D mesh from heightmap (which can be a 2D array or an image (.tif, .png, .jpg, etc), and save it in
+ the obj format.
+
+ Args:
+ heightmap (float[M,N], str): 2D array where each value represents the height. If it is a string, it is assumed
+ that is the path to a file .tif, .geotiff or an image (.png, .jpg, etc). It will be opened using the
+ `gdal` library.
+ x (float[M,N], None): 2D array where each value represents the x position (array from meshgrid). If None, it
+ will generate it automatically from the heightmap. If `heightmap` is a string, this `x` won't be taken
+ into account.
+ y (float[M,N], None): 2D array where each value represents the y position (array from meshgrid). If None, it
+ will generate it automatically from the heightmap. If `heightmap` is a string, this `y` won't be taken
+ into account.
+ depth_level (float): the depth will be the minimum depth of the heightmap minus the given depth_level.
+ filename (str, None): filename to save the mesh. If None, it won't save it.
+ show (bool): if True, it will show the mesh using `mayavi.mlab`.
+ subsample (int, None): if not None, it is the number of points to sub-sample (to smooth the heightmap using
+ the specified function)
+ interpolate_fct (str, callable): "The radial basis function, based on the radius, r, given by the norm
+ (default is Euclidean distance);
+ 'multiquadric': sqrt((r/self.epsilon)**2 + 1)
+ 'inverse': 1.0/sqrt((r/self.epsilon)**2 + 1)
+ 'gaussian': exp(-(r/self.epsilon)**2)
+ 'linear': r
+ 'cubic': r**3
+ 'quintic': r**5
+ 'thin_plate': r**2 * log(r)
+ If callable, then it must take 2 arguments (self, r). The epsilon parameter will be available as
+ self.epsilon. Other keyword arguments passed in will be available as well." [1]
+ lower_bound (int, float, None): lower bound; each value in the heightmap will be higher than or equal to
+ this bound
+ upper_bound (int, float, None): upper bound; each value in the heightmap will be lower than or equal to
+ this bound
+ dtype (np.int, np.float, None): type of the returned array for the heightmap
+ center (bool): if True, it will center the mesh
+
+ Examples:
+ import numpy as np
+
+ height = np.random.rand(100,100) # in meters
+ create3DMesh(height, show=True)
+ """
+ if isinstance(heightmap, str):
+ # load data (raster)
+ data = gdal.Open(heightmap)
+
+ gt = data.GetGeoTransform()
+ # gt is an array with:
+ # 0 = x-coordinate of the upper-left corner of the upper-left pixel
+ # 1 = width of a pixel
+ # 2 = row rotation (typically zero)
+ # 3 = y-coordinate of the of the upper-left corner of the upper-left pixel
+ # 4 = column rotation (typically zero)
+ # 5 = height of a pixel (typically negative)
+
+ # numpy array of shape: (channel, height, width)
+ #dem = data.ReadAsArray()
+
+ # get elevation values (i.e. height values) with shape (height, width)
+ band = data.GetRasterBand(1)
+ band = band.ReadAsArray()
+
+ # generate coordinates (x,y,z)
+ xres, yres = gt[1], gt[5]
+ width, height = data.RasterXSize * xres, data.RasterYSize * yres
+ xmin = gt[0] + xres * 0.5
+ xmax = xmin + width - xres * 0.5
+ ymin = gt[3] + yres * 0.5
+ ymax = ymin + height - yres * 0.5
+
+ x, y = np.arange(xmin, xmax, xres), np.arange(ymin, ymax, yres)
+ x, y = np.meshgrid(x, y)
+ z = band
+
+ # if we need to subsample, it will smooth the heightmap
+ if isinstance(subsample, int) and subsample > 0:
+ height, width = z.shape
+ idx_x = np.linspace(0, height - 1, subsample, dtype=np.int)
+ idx_y = np.linspace(0, width - 1, subsample, dtype=np.int)
+ idx_x, idx_y = np.meshgrid(idx_x, idx_y)
+ rbf = scipy.interpolate.Rbf(x[idx_x, idx_y], y[idx_x, idx_y], z[idx_x, idx_y], function=interpolate_fct)
+ # Nx, Ny = x.shape[0] / subsample, x.shape[1] / subsample
+ # rbf = Rbf(x[::Nx, ::Ny], y[::Nx, ::Ny], z[::Nx, ::Ny], function=interpolate_fct)
+ z = rbf(x, y)
+
+ # make sure the values of the heightmap are between the bounds (in-place), and is the correct type
+ if lower_bound and upper_bound:
+ np.clip(z, lower_bound, upper_bound, z)
+ elif lower_bound:
+ np.clip(z, lower_bound, z.max(), z)
+ elif upper_bound:
+ np.clip(z, z.min(), upper_bound, z)
+ if dtype:
+ z.astype(dtype)
+
+ else:
+ # check the heightmap is a 2D array
+ if not isinstance(heightmap, np.ndarray):
+ raise TypeError("Expecting a 2D numpy array")
+ if len(heightmap.shape) != 2:
+ raise ValueError("Expecting a 2D numpy array")
+
+ z = heightmap
+ if x is None or y is None:
+ height, width = z.shape
+ x, y = np.meshgrid(np.arange(width), np.arange(height))
+
+ # center the coordinates if specified
+ if center:
+ x,y = recenter([x,y])
+
+ # create lower plane
+ z0 = np.min(z) * np.ones(z.shape) - depth_level
+
+ # create left, right, front, and back planes
+ c1 = (np.vstack((x[0], x[0])), np.vstack((y[0], y[0])), np.vstack((z0[0], z[0])))
+ c2 = (np.vstack((x[-1], x[-1])), np.vstack((y[-1], y[-1])), np.vstack((z0[-1], z[-1])))
+ c3 = (np.vstack((x[:, 0], x[:, 0])), np.vstack((y[:, 0], y[:, 0])), np.vstack((z0[:, 0], z[:, 0])))
+ c4 = (np.vstack((x[:, -1], x[:, -1])), np.vstack((y[:, -1], y[:, -1])), np.vstack((z0[:, -1], z[:, -1])))
+ c = [c1, c2, c3, c4]
+
+ # createMesh([x, x] + [i[0] for i in c], [y, y] + [i[1] for i in c], [z, z0] + [i[2] for i in c],
+ # filename=filename, show=show, center=False)
+ createMesh([x, x] + [i[0] for i in c], [y, y] + [i[1] for i in c], [z, z0] + [i[2] for i in c],
+ filename=filename, show=show, center=False)
+
+
+def createURDFFromMesh(meshfile, filename, position=(0.,0.,0.), orientation=(0.,0.,0.), scale=(1.,1.,1.),
+ color=(1,1,1,1), texture=None, mass=0., inertia=(0.,0.,0.,0.,0.,0.),
+ lateral_friction=0.5, rolling_friction=0., spinning_friction=0., restitution=0.,
+ kp=None, kd=None): #, cfm=0., erf=0.):
+ """
+ Create a URDF file and insert the specified mesh inside.
+
+ Args:
+ meshfile (str): path to the mesh file
+ filename (str): filename of the urdf
+ position (float[3]): position of the mesh
+ orientation (float[3]): orientation (roll, pitch, yaw) of the mesh
+ scale (float[3]): scale factor in the x, y, z directions
+ color (float[4]): RGBA color where rgb=(0,0,0) is for black, rgb=(1,1,1) is for white, and a=1 means opaque.
+ texture (str, None): path to the texture to be applied to the object. If None, provided it will use the
+ given color.
+ mass (float): mass in kg
+ inertia (float[6]): upper/lower triangle of the inertia matrix (read from left to right, top to bottom)
+ lateral_friction (float): friction coefficient
+ rolling_friction (float): rolling friction coefficient orthogonal to contact normal
+ spinning_friction (float): spinning friction coefficient around contact normal
+ kp (float, None): contact stiffness (useful to make surfaces soft). Set it to None/-1 if not using it.
+ kd (float, None): contact damping (useful to make surfaces soft). Set it to None/-1 if not using it.
+ #cfm: constraint force mixing
+ #erp: error reduction parameter
+
+ Returns:
+ None
+
+ References:
+ - "ROS URDF Tutorial": http://wiki.ros.org/urdf/Tutorials
+ - "URDF: Link": http://wiki.ros.org/urdf/XML/link
+ - "Tutorial: Using a URDF in Gazebo": http://gazebosim.org/tutorials/?tut=ros_urdf
+ - SDF format: http://sdformat.org/spec
+ """
+ def getStr(lst):
+ return ' '.join([str(i) for i in lst])
+
+ position = getStr(position)
+ orientation = getStr(orientation)
+ color = getStr(color)
+ scale = getStr(scale)
+ name = meshfile.split('/')[-1][:-4]
+ ixx, ixy, ixz, iyy, iyz, izz = [str(i) for i in inertia]
+
+ with open(filename, 'w') as f:
+ f.write('')
+ f.write('')
+ f.write('\t')
+
+ f.write('\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ if kp is not None:
+ f.write('\t\t\t')
+ if kd is not None:
+ f.write('\t\t\t')
+ # f.write('\t\t\t')
+ # f.write('\t\t\t')
+ # f.write('\t\t\t')
+ f.write('\t\t')
+
+ f.write('\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t')
+
+ f.write('\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ if texture is not None:
+ f.write('\t\t\t\t')
+ else:
+ f.write('\t\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t')
+
+ f.write('\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t\t\t')
+ f.write('\t\t\t')
+ f.write('\t\t')
+
+ f.write('\t')
+ f.write('')
+
+
+
+def convertX3dToObj(filename, removeX3d=True):
+ """
+ Convert a .x3d into an .obj file.
+
+ Warnings: This method use the `meshlabserver` bash command. Be sure that `meshlab` is installed on the computer.
+
+ Args:
+ filename (str): path to the .x3d file
+ removeX3d (bool): True if it should remove the old .x3d file.
+
+ Returns:
+ None
+ """
+ obj_filename = filename[:-4] + '.obj'
+
+ try:
+ # convert mesh (check `meshlabserver` command for more info)
+ subprocess.call(['meshlabserver', '-i', filename, '-o', obj_filename]) # same as calling Popen(...).wait()
+
+ # replace all commas by dots
+ for line in fileinput.input(obj_filename, inplace=True):
+ line = line.replace(',', '.')
+ sys.stdout.write(line)
+
+ # remove the old .x3d file if specified
+ if removeX3d:
+ subprocess.call(['rm', filename])
+ except OSError as e:
+ if e.errno == os.errno.ENOENT:
+ raise OSError(
+ "The command `meshlabserver` is not installed on this system. Verify that meshlab is installed.")
+ else:
+ raise OSError("Error while running the command `meshlabserver`: {}".format(e))
+
+
+def convertMesh(fromFilename, toFilename, removeFile=True):
+ """
+ Convert the given file containing the original mesh to the other specified format.
+ The available formats are the ones supported by `meshlab`.
+
+ Args:
+ fromFilename (str): filename of the mesh to convert
+ toFilename (str): filename of the converted mesh
+ removeFile (bool): True if the previous file should be deleted
+
+ Returns:
+ None
+ """
+ try:
+ # convert mesh (check `meshlabserver` command for more info)
+ subprocess.call(['meshlabserver', '-i', fromFilename, '-o', toFilename]) # same as calling Popen(...).wait()
+
+ # replace all commas by dots
+ for line in fileinput.input(toFilename, inplace=True):
+ line = line.replace(',', '.')
+ sys.stdout.write(line)
+
+ # remove the old .x3d file if specified
+ if removeFile:
+ subprocess.call(['rm', fromFilename])
+ except OSError as e:
+ if e.errno == os.errno.ENOENT:
+ raise OSError(
+ "The command `meshlabserver` is not installed on this system. Verify that meshlab is installed.")
+ else:
+ raise OSError("Error while running the command `meshlabserver`: {}".format(e))
+
+
+def readObjFile(filename):
+ r"""
+ Read an .obj file and returns the whole file, as well as the list of vertices, and faces.
+
+ Args:
+ filename (str): path to the obj file
+
+ Returns:
+ list[str]: each line in the file
+ np.array[N,3]: list of vertices, where each vertex is a 3D position
+ list[list[M]]: list of faces, where each face is a list of vertex ids which composed the face. Note that the
+ first vertex id starts from 0 and not 1 like in the file.
+ """
+ data, vertices, faces = [], [], []
+
+ with open(filename) as f:
+ for i, line in enumerate(f):
+ data.append(line)
+ words = line.split()
+ if len(words) > 0:
+ if words[0] == 'v': # vertex
+ if len(words) > 3:
+ x, y, z = words[1:4]
+ vertices.append(np.array([float(x), float(y), float(z)]))
+ elif words[0] == 'f': # face
+ face = []
+ for word in words[1:]:
+ numbers = word.split('//')
+ if len(numbers) > 0:
+ face.append(int(numbers[0]) - 1)
+ faces.append(face)
+
+ vertices = np.array(vertices)
+ return data, vertices, faces
+
+
+def flipFaceNormalsInObj(filename):
+ """
+ Flip all the face normals in .obj file.
+
+ Args:
+ filename (str): path to the obj file
+ """
+ # read (load) all the file
+ with open(filename) as f:
+ data = f.readlines()
+
+ # flip the faces
+ for i in range(len(data)):
+ words = data[i].split()
+ if len(words) > 0:
+ if words[0] == 'f': # face
+ data[i] = words[0] + ' ' + words[-1] + ' ' + words[-2] + ' ' + words[-3] + '\n'
+
+ # rewrite the obj file
+ with open(filename, 'w') as f:
+ f.writelines(data)
+
+
+def flipFaceNormalsForConvexObj(filename, outward=True):
+ """
+ Flip the face normals for convex objects, and rewrite the obj file
+
+ Args:
+ filename (str): the path to the obj file
+ outward (bool): if the face normals should point outward. If False, they will be flipped such that they point
+ inward the object.
+ """
+ # read the obj file
+ data, vertices, faces = readObjFile(filename)
+
+ # compute the center of the object
+ center = np.mean(vertices, axis=0)
+ print('Center of object: {}'.format(center))
+
+ # flip the faces that points inward or outward
+ v = vertices
+ face_id = 0
+ for i in range(len(data)):
+ words = data[i].split()
+ if len(words) > 0:
+ if words[0] == 'f': # face
+ # compute the center of the face
+ face = faces[face_id]
+ face_center = np.mean([v[face[i]] for i in range(len(face))], axis=0)
+ print('Face id: {}'.format(face_id))
+ print('Face center: {}'.format(face_center))
+
+ # compute the surface vector that goes from the center of the object to the face center
+ vector = face_center - center
+
+ # compute the normal vector of the face
+ normal = np.cross( (v[face[2]] - v[face[1]]), (v[face[0]] - v[face[1]]) )
+
+ # compute the dot product between the normal and the surface vector
+ direction = np.dot(vector, normal)
+
+ print('direction: {}'.format(direction))
+
+ # flip the faces that need to be flipped
+ if (direction > 0 and not outward) or (direction < 0 and outward):
+ data[i] = words[0] + ' ' + words[-1] + ' ' + words[-2] + ' ' + words[-3] + '\n'
+
+ # increment face id
+ face_id +=1
+
+ # rewrite the obj file
+ with open(filename, 'w') as f:
+ f.writelines(data)
+
+
+def flipFaceNormalsForExpandedObj(filename, expanded_filename, outward=True, remove_expanded_file=False):
+ r"""
+ By comparing the expanded object with the original object, we can compute efficiently the normal vector to each
+ face such that it points outward. Then comparing the direction of these obtained normal vectors with the ones
+ computed for the original faces, we can correct them.
+
+ Args:
+ filename (str): the path to the original obj file
+ expanded_filename (str): the path to the expanded obj file; the file that contains the same object but which
+ has been expanded in every dimension.
+ outward (bool): if the face normals should point outward. If False, they will be flipped such that they point
+ inward the object.
+ """
+ # read the obj files
+ d1, v1, f1 = readObjFile(filename)
+ d2, v2, f2 = readObjFile(expanded_filename)
+
+ # check the size of the obj files (they have to match)
+ if len(v1) != len(v2) or len(f1) != len(f2):
+ raise ValueError("Expecting to have the same number of vertices and faces in each file: "
+ "v1={}, v2={}, f1={}, f2={}".format(len(v1), len(v2), len(f1), len(f2)))
+ if len(d1) != len(d2):
+ raise ValueError("Expecting the files to have the same size, but instead we have {} and {}".format(len(d1),
+ len(d2)))
+
+ # flip the faces that points inward or outward
+ face_id = 0
+ for i in range(len(d1)):
+ words = d1[i].split()
+ if len(words) > 0:
+ if words[0] == 'f': # face
+ # compute the center of the faces
+ face1, face2 = f1[face_id], f2[face_id]
+ face1_center = np.mean([v1[face1[i]] for i in range(len(face1))], axis=0)
+ face2_center = np.mean([v2[face2[i]] for i in range(len(face2))], axis=0)
+
+ # compute the surface vector that goes from the original face to the expanded one
+ vector = face2_center - face1_center
+
+ # compute the normal vector of the face
+ normal = np.cross((v1[face1[2]] - v1[face1[1]]), (v1[face1[0]] - v1[face1[1]]))
+
+ # compute the dot product between the normal and the surface vector
+ direction = np.dot(vector, normal)
+
+ # flip the faces that need to be flipped
+ if (direction < 0 and not outward) or (direction > 0 and outward):
+ d1[i] = words[0] + ' ' + words[-1] + ' ' + words[-2] + ' ' + words[-3] + '\n'
+
+ # increment face id
+ face_id += 1
+
+ # rewrite the obj file
+ with open(filename, 'w') as f:
+ f.writelines(d1)
+
+ # remove the expanded file
+ if remove_expanded_file:
+ os.remove(expanded_filename)
+
+
+# Test
+if __name__ == '__main__':
+
+ # 1. create 3D ellipsoid mesh (see `https://en.wikipedia.org/wiki/Ellipsoid` for more info)
+ a,b,c,n = 1., 0.5, 0.5, 50
+ #a,b,c,n = .5, .5, .5, 37
+ theta, phi = np.meshgrid(np.linspace(-np.pi/2, np.pi/2, n), np.linspace(-np.pi, np.pi, n))
+
+ x = a * np.cos(theta) * np.cos(phi)
+ y = b * np.cos(theta) * np.sin(phi)
+ z = c * np.sin(theta)
+
+ createMesh(x, y, z, show=True)
+ #createMesh(x, y, z, filename='ellipsoid.obj', show=True)
+
+ # 2. create heightmap mesh
+ height = np.random.rand(100,100) # in meters
+ createSurfMesh(height, show=True)
+
+ # 3. create right triangular prism
+ x = np.array([[-0.5,-0.5],
+ [0.5, 0.5],
+ [-0.5,-0.5],
+ [-0.5,-0.5],
+ [-0.5,0.5],
+ [0.5,-0.5],
+ [-0.5, 0.5],
+ [0.5, -0.5]])
+ y = np.array([[-0.5,0.5],
+ [-0.5,0.5],
+ [-0.5,0.5],
+ [-0.5,0.5],
+ [-0.5,-0.5],
+ [-0.5,-0.5],
+ [0.5, 0.5],
+ [0.5, 0.5]])
+ z = np.array([[0.,0.],
+ [0.,0.],
+ [1.,1.],
+ [0.,0.],
+ [0.,0.],
+ [0.,1.],
+ [0., 0.],
+ [0., 1.]])
+
+ #createMesh(x, y, z, show=True)
+ createMesh(x, y, z, filename='right_triangular_prism.obj', show=True)
+ flipFaceNormalsForConvexObj('right_triangular_prism.obj', outward=True)
+
+ exit()
+
+ # 4. create cone
+ radius, height, n = 0.5, 1., 50
+ [r, theta] = np.meshgrid((radius, 0.), np.linspace(0, 2*np.pi, n))
+ [h, theta] = np.meshgrid((0., height), np.linspace(0, 2*np.pi, n))
+ x, y, z = r * np.cos(theta), r * np.sin(theta), h
+ # close the cone at the bottom
+ [r, theta] = np.meshgrid((0., radius), np.linspace(0, 2*np.pi, n))
+ x = np.vstack((x, r * np.cos(theta)))
+ y = np.vstack((y, r * np.sin(theta)))
+ z = np.vstack((z, np.zeros(r.shape)))
+
+ createMesh(x, y, z, show=True)
+ #createMesh(x, y, z, filename='cone.obj', show=True)
+
+ # 5. create 3D heightmap
+ dx, dy, dz = 5., 5., 0.01
+ x,y = np.meshgrid(np.linspace(-dx, dx, int(2*dx)), np.linspace(-dy, dy, int(2*dy)))
+ z = np.random.rand(*x.shape) + dz
+
+ # z0 = np.zeros(x.shape)
+ #
+ # w = np.dstack((x,y,z0,z)) # 2DX x 2DY x 4
+ #
+ # c1 = (np.vstack((x[0], x[0])), np.vstack((y[0],y[0])), np.vstack((z0[0],z[0])))
+ # c2 = (np.vstack((x[-1], x[-1])), np.vstack((y[-1],y[-1])), np.vstack((z0[-1],z[-1])))
+ # c3 = (np.vstack((x[:,0], x[:,0])), np.vstack((y[:,0], y[:,0])), np.vstack((z0[:,0], z[:,0])))
+ # c4 = (np.vstack((x[:,-1], x[:,-1])), np.vstack((y[:,-1], y[:,-1])), np.vstack((z0[:,-1], z[:,-1])))
+ # c = [c1,c2,c3,c4]
+ #
+ # createMesh([x,x]+[i[0] for i in c], [y,y]+[i[1] for i in c], [z,z0]+[i[2] for i in c], show=True)
+
+ create3DMesh(z, x, y, dz, show=True)
diff --git a/pyrobolearn/utils/mocap_parser.py b/pyrobolearn/utils/mocap_parser.py
new file mode 100644
index 0000000..9ec4976
--- /dev/null
+++ b/pyrobolearn/utils/mocap_parser.py
@@ -0,0 +1,230 @@
+import numpy as np
+from scipy.interpolate import CubicSpline
+import matplotlib.pyplot as plt
+
+class MocapParser(object):
+
+ def __init__(self, filename):
+ """
+ Parser for motion capture. By default, if the data is described in Cartesian space, the x-axis should
+ be pointing in front of the human, the y-axis on his/her left, and z-axis upward.
+ :param filename:
+ """
+ self.filename = filename
+ self.num_samples = 0
+ self.joint_names = []
+ self.link_names = []
+ self.marker_names = []
+
+ self.data = self.loadFile(filename)
+
+
+ def loadFile(self, filename):
+ raise NotImplementedError("loadFile is not implemented.")
+
+ def interpolate(self, data, method='cubic', axis=-1)
+ """
+ Interpolate the Mocap data such that it is between 0 and 1, along the given axis.
+
+ :param data: mocap data
+ :param method: 'linear', 'cubic', 'hermite' interpolation
+ :param axis: The axis on which to interpolate. The length should be equal to the number of samples in the
+ mocap data
+ :return: Interpolator - function that given the time [0,1] will give the corresponding data
+ """
+ self.num_samples = data.shape[axis]
+ x = np.linspace(0., 1., self.num_samples)
+ interpolator = CubicSpline(x, self.data, axis=axis)
+ return interpolator
+
+ def getMarkerName(self, marker_idx=None):
+ if marker_idx is None:
+ return self.getMarkerNames()
+ else:
+ return self.marker_names[marker_idx]
+
+ def getMarkerNames(self):
+ return self.marker_names
+
+ def getJointName(self, joint_idx=None):
+ if joint_idx is None:
+ return self.getJointNames()
+ else:
+ return self.joint_names[joint_idx]
+
+ def getJointNames(self):
+ return self.joint_names
+
+ def getLinkName(self, link_idx=None):
+ if link_idx is None:
+ return self.getLinkNames()
+ else:
+ return self.link_names[link_idx]
+
+ def getLinkNames(self):
+ return self.link_names
+
+ def getMarkerPosition(self, marker_idx=None):
+ if marker_idx is None:
+ return self.getMarkerPositions()
+ else:
+ pass
+
+ def getMarkerPositions(self):
+ pass
+
+ def getJointPosition(self, joint_idx=None):
+ if joint_idx is None:
+ return self.getJointPositions()
+ else:
+ pass
+
+ def getJointPositions(self):
+ pass
+
+ def getJointVelocity(self, joint_idx=None):
+ if joint_idx is None:
+ return self.getJointVelocities()
+ else:
+ pass
+
+ def getJointVelocities(self):
+ pass
+
+ def getLinkPosition(self, link_idx=None):
+ if link_idx is None:
+ return self.getLinkPositions()
+ else:
+ pass
+
+ def getLinkPositions(self):
+ pass
+
+ def getLinkVelocity(self, link_idx=None):
+ if link_idx is None:
+ return self.getLinkVelocities()
+ else:
+ pass
+
+ def getLinkVelocities(self):
+ pass
+
+ def getLinkOrientation(self, link_idx=None):
+ if link_idx is None:
+ return self.getLinkOrientations()
+ else:
+ pass
+
+ def getLinkOrientations(self):
+ pass
+
+ def getLinkAngularVelocity(self, link_idx=None):
+ if link_idx is None:
+ return self.getLinkAngularVelocities()
+ else:
+ pass
+
+ def getLinkAngularVelocities(self):
+ pass
+
+
+ ## Plotting ##
+ def plot3d(self, ax=None):
+ pass
+
+ def plotJointProfile(self, ax=None, joint_idx=None, pos=True, vel=True, acc=True):
+ pass
+
+ def plotLinkProfile(self, ax=None, link_idx=None, pos=True, vel=True, acc=True, wrt='world'):
+ pass
+
+ def plotMarkerProfile(self, ax=None, link_idx=None, pos=True, vel=True, acc=True, wrt='world'):
+ pass
+
+ def animate3d(self, ax=None, title=None):
+ pass
+
+
+
+
+from amcparser.skeleton import Skeleton
+from amcparser.motion import SkelMotion
+
+class CMUMocapParser(MocapParser):
+
+ def __init__(self, skeleton_filename, motion_filename, skeleton_scale=1.0):
+ super(CMUMocapParser, self).__init__(motion_filename)
+ self.joint_names = ['head', 'upperneck', 'lowerneck', 'upperback', 'thorax', 'lowerback', 'root', # Spine
+ 'rclavicle', 'rhumerus', 'rradius', 'rwrist', 'rhand', 'rthumb', 'rfingers', # Right arm
+ 'lclavicle', 'lhumerus', 'lradius', 'lwrist', 'lhand', 'lthumb', 'lfingers', # Left arm
+ 'rhipjoint', 'rfemur', 'rtibia', 'rfoot', 'rtoes', # Right leg
+ 'lhipjoint', 'lfemur', 'ltibia', 'lfoot', 'ltoes'] # Left leg
+ self.link_names = self.joint_names
+ self.marker_names = self.joint_names
+ self.base_name = 'root'
+
+ # Load skeleton
+ self.skeleton = Skeleton(skeleton_filename, scale=skeleton_scale)
+
+ def loadFile(self, filename, framerate=120.):
+ self.skeleton_motion = SkelMotion(self.skeleton, filename, (1./framerate))
+ # compute trajectories
+ #self.data = self.skeleton_motion.traverse(bone, start, end)
+ self.data = self.skeleton_motion.traverse(None, 0, -1)
+ # make sure that given axis
+
+ def animate3d(self, ax=None, title=None):
+ if ax is None:
+ fig = plt.figure()
+ ax = fig.gca(projection='3d')
+
+ # Rescaling such that the skeleton it is in the right proportion and at the middle
+ xmin, xmax = X[..., 2].min(), X[..., 2].max()
+ ymin, ymax = X[..., 0].min(), X[..., 0].max()
+ zmin, zmax = X[..., 1].min(), X[..., 1].max()
+ x_len, y_len, z_len = (xmax - xmin), (ymax - ymin), (zmax - zmin)
+ max_len = max([x_len, y_len, z_len])
+ xmin, xmax = xmin + (x_len - max_len) / 2., xmin + (x_len + max_len) / 2.
+ ymin, ymax = ymin + (y_len - max_len) / 2., ymin + (y_len + max_len) / 2.
+ zmin, zmax = zmin + (z_len - max_len) / 2., zmin + (z_len + max_len) / 2.
+
+ # Plot trajectories
+ x, y, z = skel.bones['rhand'].xyz_data.T
+ T = len(x)
+
+ def init():
+ ax.set_title('movement')
+ ax.set_xlabel('x')
+ ax.set_xlim(xmin, xmax)
+ ax.set_ylabel('y')
+ ax.set_ylim(ymin, ymax)
+ ax.set_zlabel('z')
+ ax.set_zlim(zmin, zmax)
+ # ax.scatter(x[0], y[0], z[0], marker='o')
+ return fig,
+
+ def animate(i):
+ # ax.view_init(elev=10., azim=i)
+ ax.scatter(x[i], y[i], z[i], marker='o')
+ return fig,
+
+ def animate_skeleton(i):
+ ax.clear()
+ init()
+ # ax.scatter(X[:,2,i], X[:,0,i], X[:,1,i], marker='o', c='b')
+ for d in [X_TO, X_RA, X_LA, X_RL, X_LL]:
+ ax.plot(d[:, i, 2], d[:, i, 0], d[:, i, 1], marker='o', c='b')
+ return [fig] # fig,
+
+ # Animate
+ # anim = animation.FuncAnimation(fig, animate, init_func=init,
+ # frames=T, interval=20, blit=True)
+ anim = animation.FuncAnimation(fig, animate_skeleton, init_func=init,
+ frames=T, interval=20, blit=False)
+
+ plt.show()
+
+
+# Test
+if __name__ == "__main__":
+ pass
\ No newline at end of file
diff --git a/pyrobolearn/utils/multiprocessing_utils.py b/pyrobolearn/utils/multiprocessing_utils.py
new file mode 100644
index 0000000..fc08808
--- /dev/null
+++ b/pyrobolearn/utils/multiprocessing_utils.py
@@ -0,0 +1,2 @@
+# pathos
+# openmpi
diff --git a/pyrobolearn/utils/orientation.py b/pyrobolearn/utils/orientation.py
new file mode 100644
index 0000000..3c0b305
--- /dev/null
+++ b/pyrobolearn/utils/orientation.py
@@ -0,0 +1,333 @@
+# utils code to transform orientation expressed in different forms
+# This includes rotation matrices, euler angles (RPY), axis-angle, and quaternions
+
+import numpy as np
+import quaternion
+import sympy
+from collections import Iterable
+
+from converter import QuaternionNumpyConverter
+
+
+def getMatrixFromAxisAngle(axis, angle):
+ x, y, z = axis
+ a = angle
+ c, s = np.cos(a), np.sin(a)
+ c1 = 1 - c
+ R = np.array([[x ** 2 * c1 + c, x * y * c1 - z * s, x * z * c1 + y * s],
+ [x * y * c1 + z * s, y ** 2 * c1 + c, y * z * c1 - x * s],
+ [x * z * c1 - y * s, y * z * c1 + x * s, z ** 2 * c1 + c]])
+ return R
+
+
+def getSymbolicMatrixFromAxisAngle(axis, angle):
+ x, y, z = axis
+ a = angle
+ c, s = sympy.cos(a), sympy.sin(a)
+ c1 = 1 - c
+ R = np.array([[x**2 * c1 + c, x * y * c1 - z * s, x * z * c1 + y * s],
+ [x * y * c1 + z * s, y**2 * c1 + c, y * z * c1 - x * s],
+ [x * z * c1 - y * s, y * z * c1 + x * s, z**2 * c1 + c]])
+ return R
+
+
+def getAxisAngleFromMatrix(R):
+ angle = np.arccos((R[0, 0] + R[1, 1] + R[2, 2] - 1) / 2.)
+ axis = 1. / (2. * np.sin(angle)) * np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
+ return angle, axis
+
+
+def getSymbolicAxisAngleFromMatrix(R):
+ angle = sympy.acos((R[0, 0] + R[1, 1] + R[2, 2] - 1) / 2.)
+ axis = 1. / (2. * sympy.sin(angle)) * np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
+ return angle, axis
+
+
+def getQuaternionFromAxisAngle(axis, angle, convert_to_quat=False, convention='xyzw'):
+ w = np.cos(angle / 2.)
+ x, y, z = np.sin(angle / 2.) * axis
+ if convert_to_quat:
+ return quaternion.quaternion(w, x, y, z)
+ else:
+ if convention == 'xyzw':
+ return np.array([x, y, z, w])
+ elif convention == 'wxyz':
+ return np.array([w, x, y, z])
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+
+
+def getSymbolicQuaternionFromAxisAngle(axis, angle, convention='xyzw'):
+ w = sympy.cos(angle / 2.)
+ x, y, z = sympy.sin(angle / 2.) * axis
+ if convention == 'xyzw':
+ return np.array([x, y, z, w])
+ elif convention == 'wxyz':
+ return np.array([w, x, y, z])
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+
+
+def getRPYFromMatrix(R):
+ r = np.arctan2(R[1, 0], R[0, 0])
+ p = np.arctan2(-R[2, 0], np.sqrt(R[2, 1]**2 + R[2, 2]**2))
+ y = np.arctan2(R[2, 1], R[2, 2])
+ return np.array([r, p, y])
+
+
+def getSymbolicRPYFromMatrix(R):
+ r = sympy.atan2(R[1, 0], R[0, 0])
+ p = sympy.atan2(-R[2, 0], sympy.sqrt(R[2, 1] ** 2 + R[2, 2] ** 2))
+ y = sympy.atan2(R[2, 1], R[2, 2])
+ return np.array([r, p, y])
+
+
+def getMatrixFromRPY(rpy):
+ cr, cp, cy = [np.cos(i) for i in rpy]
+ sr, sp, sy = [np.sin(i) for i in rpy]
+ R = np.array([[cy*cp, cy*sp*sr - sy*cr, cy*sp*cr + sy*sr],
+ [sy*cp, sy*sp*sr + cy*cr, sy*sp*cr - cy*sr],
+ [-sp, cp*sr, cp*cr]])
+ return R
+
+
+def getSymbolicMatrixFromRPY(rpy):
+ cr, cp, cy = [sympy.cos(i) for i in rpy]
+ sr, sp, sy = [sympy.sin(i) for i in rpy]
+ R = np.array([[cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
+ [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
+ [-sp, cp * sr, cp * cr]])
+ return R
+
+
+def getQuaternionFromMatrix(R, convert_to_quat=False, convention='xyzw'):
+ w = 1./2 * np.sqrt(R[0, 0] + R[1, 1] + R[2, 2] + 1)
+ x, y, z = 1./2 * np.array([np.sign(R[2, 1] - R[1, 2]) * np.sqrt(R[0, 0] - R[1, 1] - R[2, 2] + 1),
+ np.sign(R[0, 2] - R[2, 0]) * np.sqrt(R[1, 1] - R[2, 2] - R[0, 0] + 1),
+ np.sign(R[1, 0] - R[0, 1]) * np.sqrt(R[2, 2] - R[0, 0] - R[1, 1] + 1)])
+ if convert_to_quat:
+ return quaternion.quaternion(w, x, y, z)
+ else:
+ if convention == 'xyzw':
+ return np.array([x, y, z, w])
+ elif convention == 'wxyz':
+ return np.array([w, x, y, z])
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+
+
+def getSymbolicQuaternionFromMatrix(R, convention='xyzw'):
+ w = 1. / 2 * sympy.sqrt(R[0, 0] + R[1, 1] + R[2, 2] + 1)
+ x, y, z = 1. / 2 * np.array([sympy.sign(R[2, 1] - R[1, 2]) * sympy.sqrt(R[0, 0] - R[1, 1] - R[2, 2] + 1),
+ sympy.sign(R[0, 2] - R[2, 0]) * sympy.sqrt(R[1, 1] - R[2, 2] - R[0, 0] + 1),
+ sympy.sign(R[1, 0] - R[0, 1]) * sympy.sqrt(R[2, 2] - R[0, 0] - R[1, 1] + 1)])
+ if convention == 'xyzw':
+ return np.array([x, y, z, w])
+ elif convention == 'wxyz':
+ return np.array([w, x, y, z])
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+
+
+def getMatrixFromQuaternion(q, convention='xyzw'):
+ if isinstance(q, quaternion.quaternion):
+ x, y, z, w = q.x, q.y, q.z, q.w
+ elif isinstance(q, Iterable):
+ if convention == 'xyzw':
+ x, y, z, w = q
+ elif convention == 'wxyz':
+ w, x, y, z = q
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+ else:
+ raise TypeError
+ R = np.array([[2 * (w**2 + x**2) - 1, 2 * (x*y - w*z), 2 * (x*z + w*y)],
+ [2 * (x*y + w*z), 2 * (w**2 + y**2) - 1, 2*(y*z - w*x)],
+ [2 * (x*z - w*y), 2 * (y*z + w*x), 2 * (w**2 + z**2) - 1]])
+ return R
+
+
+def getSymbolicMatrixFromQuaternion(q, convention='xyzw'):
+ return getMatrixFromQuaternion(q, convention=convention)
+
+
+def skew(vector):
+ r"""
+ Return the skew-symmetric matrix of the given vector, which allows to represents the cross product between the
+ given vector and another vector, as the multiplication of the returned skew-symmetric matrix with the other
+ vector.
+
+ The skew-symmetric matrix from a 3D vector :math:`v=[x,y,z]` is given by:
+
+ .. math::
+
+ S(v) = \left[ \begin{array}{ccc} 0 & -z & y \\ z & 0 & -x \\ -y & x & 0 \\ \end{array} \right]
+
+ It can be shown [2] that: :math:`\dot{R}(t) = \omega(t) \times R(t) = S(\omega(t)) R(t)`, where :math:`R(t)` is
+ a rotation matrix that varies as time :math:`t` goes, :math:`\omega(t)` is the angular velocity vector of frame
+ :math:`R(t) with respect to the reference frame at time :math:`t`, and :math:`S(.)` is the skew operation that
+ returns the skew-symmetric matrix from the given vector.
+
+ Args:
+ vector (np.array[3]): 3D vector
+
+ Returns:
+ np.array[3,3]: skew-symmetric matrix
+
+ References:
+ [1] Wikipedia: https://en.wikipedia.org/wiki/Skew-symmetric_matrix#Cross_product
+ [2] "Robotics: Modelling, Planning and Control" (sec 3.1.1), by Siciliano et al., 2010
+ """
+ x, y, z = vector
+ return np.array([[0., -z, y],
+ [z, 0., -x],
+ [-y, x, 0.]])
+
+
+def RotX(angle):
+ """
+ Return the rotation matrix around the x-axis by the given angle.
+
+ Args:
+ angle (float): angle in radians
+
+ Returns:
+ np.array[3,3]: rotation matrix around the x-axis
+ """
+ c, s = np.cos(angle), np.sin(angle)
+ return np.array([[1., 0., 0.],
+ [0., c, -s],
+ [0., s, c]])
+
+
+def RotY(angle):
+ """
+ Return the rotation matrix around the y-axis by the given angle.
+
+ Args:
+ angle (float): angle in radians
+
+ Returns:
+ np.array[3,3]: rotation matrix around the y-axis
+ """
+ c, s = np.cos(angle), np.sin(angle)
+ return np.array([[c, 0., s],
+ [0., 1., 0.],
+ [-s, 0, c]])
+
+
+def RotZ(angle):
+ """
+ Return the rotation matrix around the z-axis by the given angle.
+
+ Args:
+ angle (float): angle in radians
+
+ Returns:
+ np.array[3,3]: rotation matrix around the z-axis
+ """
+ c, s = np.cos(angle), np.sin(angle)
+ return np.array([[c, -s, 0.],
+ [s, c, 0.],
+ [0., 0., 1.]])
+
+
+###############
+# Quaternions #
+###############
+
+quat_converter = QuaternionNumpyConverter(convention=1)
+
+
+def getQuaternionInverse(q, convention='xyzw'):
+ if isinstance(q, quaternion.quaternion):
+ return q.inverse()
+ elif isinstance(q, Iterable):
+ if convention == 'xyzw':
+ x, y, z, w = q
+ return np.array([-x, -y, -z, w])
+ elif convention == 'wxyz':
+ w, x, y, z = q
+ return np.array([w, -x, -y, -z])
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+ else:
+ raise TypeError
+
+
+def getQuaternionProduct(q1, q2, convention='xyzw'):
+ if type(q1) != type(q2):
+ raise TypeError("Expecting q1 and q2 to be of the same type")
+ if isinstance(q1, quaternion.quaternion):
+ return q1*q2
+ elif isinstance(q1, Iterable):
+ if convention == 'xyzw':
+ x1, y1, z1, w1 = q1
+ x2, y2, z2, w2 = q2
+ v1, v2 = np.array([x1, y1, z1]), np.array([x2, y2, z2])
+ v = w1 * v2 + w2 * v1 + np.cross(v1, v2)
+ w = w1 * w2 - v1.dot(v2)
+ return np.array([v[0], v[1], v[2], w])
+ elif convention == 'wxyz':
+ w1, x1, y1, z1 = q1
+ w2, x2, y2, z2 = q2
+ v1, v2 = np.array([x1, y1, z1]), np.array([x2, y2, z2])
+ v = w1 * v2 + w2 * v1 + np.cross(v1, v2)
+ w = w1 * w2 - v1.dot(v2)
+ return np.array([w, v[0], v[1], v[2]])
+ else:
+ raise NotImplementedError("Asking for a convention that has not been implemented")
+ else:
+ raise TypeError
+
+
+def logarithm_map(q):
+ r"""
+ Apply the logarithm map to a quaternion; :math:`log : S^3 \rightarrow R^3`.
+
+ Args:
+ q (float[4]): quaternion
+
+ Returns:
+ float[3]: resulting 3d vector
+ """
+ q = quat_converter.convertTo(q)
+ v, u = q.w, np.array([q.x, q.y, q.z])
+
+ zero = np.zeros(3)
+ if np.allclose(u, zero):
+ return zero
+ return np.arccos(v) * u / np.linalg.norm(u)
+
+
+def exponential_map(r):
+ r"""
+ Apply the exponential map to a 3d vector representing an orientation; :math:`exp : R^3 \rightarrow S^3`
+
+ Args:
+ r (float[3]): 3d vector
+
+ Returns:
+ float[4]: quaternion
+ """
+ if np.allclose(r, np.zeros(3)):
+ return quaternion.quaternion(1, 0, 0, 0)
+ r_ = np.linalg.norm(r)
+ x, y, z = np.sin(r_) * r / r_
+ return quaternion.quaternion(np.cos(r_), x, y, z)
+
+
+def angular_velocity_from_quaternion(q1, q2):
+ """
+ Convert the difference between 2 quaternions using the logarithm map.
+
+ Args:
+ q1: first (desired) quaternion
+ q2: second (current) quaternion
+
+ Returns:
+ float[3]: angular velocity (angular error in :math:`R^3`)
+ """
+ q1 = quat_converter.convertTo(q1)
+ q2 = quat_converter.convertTo(q2)
+ return 2 * logarithm_map(q1 * q2)
diff --git a/pyrobolearn/utils/plot.py b/pyrobolearn/utils/plot.py
new file mode 100644
index 0000000..367ba19
--- /dev/null
+++ b/pyrobolearn/utils/plot.py
@@ -0,0 +1,20 @@
+# Matplotlib
+# Check also Visdom: https://github.com/facebookresearch/visdom
+
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.patches import FancyArrowPatch
+
+
+class Arrow3D(FancyArrowPatch):
+ r"""This class allows to draw a 3D arrow"""
+
+ def __init__(self, xs, ys, zs, *args, **kwargs):
+ FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs)
+ self._verts3d = xs, ys, zs
+
+ def draw(self, renderer):
+ xs3d, ys3d, zs3d = self._verts3d
+ xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
+ self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
+ FancyArrowPatch.draw(self, renderer)
\ No newline at end of file
diff --git a/pyrobolearn/utils/pose_generator.py b/pyrobolearn/utils/pose_generator.py
new file mode 100644
index 0000000..e6418fa
--- /dev/null
+++ b/pyrobolearn/utils/pose_generator.py
@@ -0,0 +1,40 @@
+# This file defines the `PoseGenerator` class which generates possible or plausible poses for a robot.
+# A pose is defined as:
+# - the link positions/orientations, and the base position/orientation.
+# - the joint positions, and the base position/orientation.
+
+
+class PoseGenerator(object):
+
+ def __init__(self, robot):
+ self.robot = robot
+
+ def generate_uniform_random_pose(self, jnts=None):
+ pass
+
+ def generate_gaussian_random_pose(self, jnts=None):
+ """
+ Put a gaussian distribution with the mean sets to jnt initial configuration, and the 2 times
+ the standard deviation sets to ...
+ :param jnts:
+ :return:
+ """
+ pass
+
+ def generate_plausible_pose(self, model, jnts=None):
+ """
+ Given a trained learning model (for instance a GAN or VAE), it generates a plausible pose of the robot.
+ :param model: learning model
+ :param jnts:
+ :return:
+ """
+ pass
+
+ def generate_random_pose(self, generator, jnts=None):
+ """
+ Based on the given distribution generator, it generates a pose.
+ :param generator:
+ :param jnts:
+ :return:
+ """
+ pass
\ No newline at end of file