From 7911bbd310a86b65c4ea2bb7387ac6592f4dc7b8 Mon Sep 17 00:00:00 2001 From: Brian Delhaisse Date: Sat, 16 Mar 2019 01:14:46 +0100 Subject: [PATCH] add utils and simulators --- pyrobolearn/__init__.py | 97 + pyrobolearn/simulators/__init__.py | 19 + pyrobolearn/simulators/bullet.py | 3411 +++++++++++++++++ pyrobolearn/simulators/bullet_ros.py | 52 + pyrobolearn/simulators/flex.py | 38 + pyrobolearn/simulators/gazebo-ros.py | 510 +++ pyrobolearn/simulators/gazebo.py | 39 + pyrobolearn/simulators/mujoco.py | 48 + pyrobolearn/simulators/opensim.py | 38 + pyrobolearn/simulators/ros.py | 63 + pyrobolearn/simulators/ros_rbdl.py | 52 + pyrobolearn/simulators/simulator.py | 403 ++ pyrobolearn/simulators/simureal.py | 62 + pyrobolearn/utils/__init__.py | 89 + pyrobolearn/utils/bullet_utils.py | 172 + pyrobolearn/utils/cmu_mocap_parser.py | 0 pyrobolearn/utils/converter.py | 484 +++ pyrobolearn/utils/data_structures/__init__.py | 8 + pyrobolearn/utils/data_structures/graph.py | 51 + .../utils/data_structures/orderedset.py | 957 +++++ pyrobolearn/utils/distributions.py | 105 + pyrobolearn/utils/gazebo-ros/spawn_model.py | 317 ++ pyrobolearn/utils/heightmap_generator.py | 424 ++ pyrobolearn/utils/human_kinematic.py | 9 + pyrobolearn/utils/interpolator.py | 110 + pyrobolearn/utils/math_utils.py | 90 + pyrobolearn/utils/mesh.py | 714 ++++ pyrobolearn/utils/mocap_parser.py | 230 ++ pyrobolearn/utils/multiprocessing_utils.py | 2 + pyrobolearn/utils/orientation.py | 333 ++ pyrobolearn/utils/plot.py | 20 + pyrobolearn/utils/pose_generator.py | 40 + 32 files changed, 8987 insertions(+) create mode 100644 pyrobolearn/__init__.py create mode 100644 pyrobolearn/simulators/__init__.py create mode 100644 pyrobolearn/simulators/bullet.py create mode 100644 pyrobolearn/simulators/bullet_ros.py create mode 100644 pyrobolearn/simulators/flex.py create mode 100644 pyrobolearn/simulators/gazebo-ros.py create mode 100644 pyrobolearn/simulators/gazebo.py create mode 100644 pyrobolearn/simulators/mujoco.py create mode 100644 pyrobolearn/simulators/opensim.py create mode 100644 pyrobolearn/simulators/ros.py create mode 100644 pyrobolearn/simulators/ros_rbdl.py create mode 100644 pyrobolearn/simulators/simulator.py create mode 100644 pyrobolearn/simulators/simureal.py create mode 100644 pyrobolearn/utils/__init__.py create mode 100644 pyrobolearn/utils/bullet_utils.py create mode 100644 pyrobolearn/utils/cmu_mocap_parser.py create mode 100644 pyrobolearn/utils/converter.py create mode 100644 pyrobolearn/utils/data_structures/__init__.py create mode 100644 pyrobolearn/utils/data_structures/graph.py create mode 100644 pyrobolearn/utils/data_structures/orderedset.py create mode 100644 pyrobolearn/utils/distributions.py create mode 100755 pyrobolearn/utils/gazebo-ros/spawn_model.py create mode 100644 pyrobolearn/utils/heightmap_generator.py create mode 100644 pyrobolearn/utils/human_kinematic.py create mode 100644 pyrobolearn/utils/interpolator.py create mode 100644 pyrobolearn/utils/math_utils.py create mode 100644 pyrobolearn/utils/mesh.py create mode 100644 pyrobolearn/utils/mocap_parser.py create mode 100644 pyrobolearn/utils/multiprocessing_utils.py create mode 100644 pyrobolearn/utils/orientation.py create mode 100644 pyrobolearn/utils/plot.py create mode 100644 pyrobolearn/utils/pose_generator.py 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