Skip to content

Code Reference

mkoroknai edited this page Apr 18, 2024 · 16 revisions

Classes

classes/mujoco_display.py

class Display

Description:

Base class for passive and active simulation. Manages window, cameras and MuJoCo.

__init__(self, xml_file_name, graphics_step, virt_parsers: list = None, mocap_parsers: list = None, connect_to_optitrack=True):

Inputs:

  • xml_file_name: string, file name of the model to be loaded
  • graphics_step: float, time between consecutive graphics frames
  • virt_parsers: list of MovingObject parser methods
  • mocap_parsers: list of MocapObject parser methods
  • connect_to_optitrack: boolean, whether to connect to Motive or not

Description:

Initializes member variables, glfw library, cameras, MuJoCo library, loads the mujoco model from xml, and looks for objects in the model that follow the naming convention using the parser methods provided. The glfw library is the window manager. The cameras are MuJoCo's MjvCamera class. There is one main camera, and one "on board" camera, that can be switched from vehicle to vehicle. The on board camera is equipped with a low pass filter for the azimuth angle, and one for the elevation angle to get rid of the high frequency judder that comes from Motive.

init_glfw(self):

Description:

Initializes the glfw library. Binds methods to window events such as mouse scroll, mouse button, cursor position and keyboard.

init_cams(self):

Description:

Initializes the main camera and the on board camera. Two pairs of LiveLFilter's (see in util/mujoco_helper.py) low-pass filters are also initialized for the on board camera. One pair for the azimuth angle and one pair for the elevation angle. The reason for requiring two filters for each angle is explained in the description of update_onboard_cam in util/mujoco_helper.py.

set_cam_position(self, azimuth: float, elevation: float, lookat: list, distance: float):

Inputs:

  • azimuth: float, rotation about the vertical axis (yaw) in degrees
  • elevation: float, rotation about the horizontal axis perpendicular to the view direction (pitch) in degrees
  • lookat: 3D point for the camera to look at (a list of 3 floats)
  • distance: float, the distance from the lookat point

Description:

Sets the attributes of the main camera (not the on board camera).

load_model(self, xml_file_name):

Inputs:

  • xml_file_name: string, file name of the model to be loaded

Description:

Saves the xml file name in a member variable, loads the model, creates MjData from MjModel, and initializes MjvScene, and MjrContext.

reload_model(self, xml_file_name, drone_names_in_motive = None):

Inputs:

  • xml_file_name: string, file name of the model to be loaded
  • drone_names_in_motive: list of strings, drone names that are being tracked in Motive

Description:

Calls load_model(), and then checks for drones in the model that follow the naming convention (see naming_convention_in_xml.txt) and creates a Drone or DroneMocap instance for each. Then sets the name_in_motive attribute of the mocap drones based on drone_names_in_motive list.

get_MovingObject_by_name_in_xml(self, name):

Inputs:

  • name: string, the xml name of the searched MovingObject

Description:

Returns the MovingObject instance whose xml name matches the name provided in the input.

get_MocapObject_by_name_in_xml(self, name):

Inputs:

  • name: string, the xml name of the searched MocapObject

Description:

Returns the MocapObject instance whose xml name matches the name provided in the input.

glfw_window_should_close(self):

Description:

Makes the glfw.window_should_close() method public.

set_key_{key}_callback(self, callback_function):

Inputs:

  • callback_function: a callable that needs to be called when a key event happens

Description:

Binds a method to key events, to make these events public.

mouse_button_callback(self, window, button, action, mods):

Inputs are those that are required by glfw.

Description:

Saves cursor position and button state in member variables when either mouse buttons are pressed.

mouse_move_callback(self, window, xpos, ypos):

Inputs are those that are required by glfw.

Description:

If left mouse button is pressed, rotates the camera about the lookat point when mouse is moving. If right button is pressed, moves the lookat point when mouse is moving.

calc_dxdy(self, window):

Inputs:

  • window: a glfw window in which cursor displacement needs to be calculated

Description:

Calculates the cursor displacement since the last known cursor position. Helper function for mouse_move_callback()

zoom(self, window, x, y):

Inputs are those that are required by glfw.

Description:

Is bound to mouse scroll callback. Moves the camera closer or further to/from the lookat point on scrolling.

key_callback(self, window, key, scancode, action, mods):

Inputs are those that are required by glfw.

Description:

Manages keyboard events. Some events are processed in this method like TAB press and SPACE press. Other events are made public by calling a callback function like key_b_callback set previously from the outside.

set_title(self, window_title):

Inputs:

  • window_title: new title for the window

Description:

Sets a new title for the window.

append_title(self, text):

Inputs:

  • text: string, whatever needs to be added to the title

Description:

Appends a string to the window title. Like when video is being recorded to let the user know what's going on.

reset_title(self):

Description:

Resets the window title to the original title, or the newest title if the original has been changed. Which means that the appended text is removed.

connect_to_Optitrack(self):

Description:

Tries to connect to Motive server. Unfortunately, if Motive is not streaming, it freezes at the line motioncapture.MotionCaptureOptitrack.

change_cam(self):

Description:

Changes active camera to be the on board camera if it's currently the main camera, and vica versa.

save_video(self, image_list, width, height):

Inputs:

  • image_list: the list of frames that have been saved during video recording
  • width: the width of the images
  • height: the height of the images

Description:

Writes the saved frames as mp4 video to hard drive at self.video_save_folder using OpenCV's VideoWriter. If the folder does not exist, it creates it. Since the recorded frames are 1D arrays, the method converts them to the desired shape OpenCV requires.

save_video_background(self):

Description:

Runs save_video() on a different thread with a copy of image_list, and sets image_list to an empty list, so that the next video recording can be started even before the saving process has finished.

set_vehicle_names(self):

Description:

If there are any mocap vehicles, it creates a small pop-up window where the user can set the mocap vehicles' name_in_motive.

classes/active_simulation.py

class ActiveSimulator(Display)

Description:

Child of Display. Handles the simulation the graphics, video recording and will handle data logging in the future.

__init__(self, xml_file_name, video_intervals, control_step, graphics_step,
               virt_parsers: list = None, mocap_parsers: list = None, connect_to_optitrack=True):

Inputs:

  • xml_file_name: string, file name of the model to be loaded
  • video_intervals: list of floats, with even number of elements. Each pair represents a time interval in seconds when video should be recorded. 0 seconds is when simulation starts. Set it to None if no video needs to be recorded.
  • control_step: float, in seconds, how frequently control is updated
  • graphics_step: float, in seconds, how frequently graphics is updated
  • virt_parsers: list of MovingObject parser methods
  • mocap_parsers: list of MocapObject parser methods
  • connect_to_optitrack: boolean, whether to connect to Motive or not

Description:

Calls the constructor of Display and saves the inputs as member variables.

update(self):

Description:

Should be called in an infinite loop. If connected to Motive, it checks for data coming in based on which it updates the mocap vehicles. It updates simulated vehicles, steps the physics, updates the window displaying the graphics, and manages video recording.

update_(self):

Description:

Should be called in an infinite loop. This method runs the simulation without graphics, so it runs as fast as it can. It does not update mocap objects since it does not run in real time.

manage_video_recording():

Description:

It is called on every update. It checks whether frames should be saved based on video_intervals. When the end of an interval is reached, it writes the saved frames as an .mp4 to hard disk on a different thread. The location of the saved video is printed to the terminal.

classes/drone.py

class SPIN_DIR(Enum)

Description:

The direction in which a propeller of a drone spins. Clockwise is 1 and counter clockwise is -1 so that the speed of the spin can be multiplied by it.

class CRAZYFLIE_PROP(Enum)

Description:

Properties of the crazyflie. OFFSET is the offset of the propellers in the x and y direction from the center of mass of the drone. OFFSET_Z is the offset of the propellers in the z direction. MOTOR_PARAM is the gear parameter of the crazyflie actuators. MAX_THRUST specifies the upper limit of the actuator control range.

class BUMBLEBEE_PROP(Enum):

Description:

Properties of Bumblebee. OFFSET_X1 is the distance of two propellers from the center of the Bumblebee in the negative direction of the Bumblebee's x axis. OFFSET_X2 is the distance of the other two propellers from the center of the Bumblebee in the positive direction of the Bumblebee's x axis. OFFSET_Y is the same on the y axis in both directions. OFFSET_Z is how far up the propellers are from the center of the Bumblebee.

class Drone

Description: A virtually simulated drone using the MuJoCo library. It manages the data that corresponds to this drone in the MjData.

__init__(self, model: mujoco.MjModel, data: mujoco.MjData, name_in_xml):

Inputs:

  • model: the MjModel loaded by the simulator
  • data: MjData created from the loaded MjModel
  • name_in_xml: the name that follows the naming convention of this drone in the xml

Description:

Saves a view of the inputs in member variables. Looks for its free joint and the joint of the propellers in the data input. Saves a view of its qpos and the qpos of each propellers as member variables. Also saves a view of its control from the data input. Makes a copy of the initial propeller angle, so that the spin can be updated such that the physics does not affect it. Saves a view of its body.xquat, free_joint.qvel, and data.sensor that corresponds to this drone. These are necessary to make the methods much more readable, because in MjData, everything is stacked on top of each other, and it's a mess.

update(self, i, control_step):

Inputs:

  • i: integer, the incremented loop variable
  • control_step: how often new control is computed in simulation time

Description:

This needs to be called in every simulation control step. If the drone has a trajectory, it gets evaluated, and it's output is passed to the drone's controllers that compute it's control. Then the new control corresponding to this drone is set in MjData.

spin_propellers(self, angle_step):

Inputs:

  • angle_step: float, by how much the propeller needs to be rotated in each step

Description:

It updates the propeller angles that have been copied in the constructor each with a small difference so that it is a bit more visually convincing, and updates the joint angle of each of the four propellers.

fake_propeller_spin(self, control_step , speed = 10):

Inputs:

  • control_step: float, in seconds, how frequently control is updated
  • speed: float, the speed of the spinning

Description:

Spins the drone's propellers if they are above 0.1m, and stops the propellers if they are below.

Inputs:

  • joint_names: a list of names of all joints in the model (can be generated by mujoco_helper.get_joint_name_list())

Returns a list of virtually simulated Drone/DroneHooked instances.

Description:

Looks for drones that follow the naming convention in the list of joints. If one is found it checks whether it's a crazyflie or a bumblebee, and whether it has a hook. Then it creates a corresponding Drone or DroneHooked instance, and places it in a list which it returns after it went through the joint list.

find_hook_for_drone(names, drone_name): (static method)

Inputs:

  • names: list of joint names in the model
  • drone_name: name of the drone that needs a hook

Returns the name of the hook in the xml, or None if it wasn't found.

class DroneHooked(Drone)

Description:

Child class of Drone. Extends Drone class with a hanging hook for picking up objects. It manages the data that corresponds to this drone in the MjData.

__init__(self, model: mujoco.MjModel, data: mujoco.MjData, name_in_xml):

Inputs are the same as those of Drone's constructor.

Description:

In addition to calling Drone's constructor, it saves the hook_name_in_xml and mines the hook qpos, qvel and rod length from MjData.

update(self, i, control_step):

Description:

Same as Drone's update.

class DroneMocap

Description:

Not simulated, only mirrors a real drone in the simulation that's being tracked by Optitrack or other motion capture systems. It manages the data that corresponds to this drone in the MjData. MuJoCo provides mocap bodies, but since those bodies cannot have child bodies, to immitate propeller spinning, each propeller needs to be defined as a separate body in the xml. If the xml is generated by util/xml_generator.py, the drones and the propellers will follow the naming convention.

__init__(self, model: mujoco.MjModel, data: mujoco.MjData, mocapid, name_in_xml, name_in_motive):

Inputs:

  • model: the loaded MjModel
  • data: the MjData created from MjModel
  • name_in_xml: the name of the drone in the xml
  • name_in_motive: to initialize drone with a name that's being streamed from Motive

Description:

Saves views of the inputs, and creates a PropellerMocap instance for each propeller.

update(self, pos, quat):

Description:

Updates the position and orientation of the drone and the propellers.

Class DroneMocapHooked(DroneMocap)

Description:

A mocap drone that has a hook hanging on a rod for picking up objects. Child class of DroneMocap. It manages the data that corresponds to this drone in the MjData. Again, since mocap bodies can't have child bodies in MuJoCo, the hook needs to be a separate body like the propellers. If it were only geoms in the drone's body, it would be glued to it, and couldn't be moving like a pendulum.

__init__(self, model: mujoco.MjModel, data: mujoco.MjData, drone_mocapid, name_in_xml, name_in_motive):

Inputs are the same as DroneMocap.

Description:

In addition to DroneMocap's constructor, it saves hook_name_in_xml in a member variable, and searches for the mocapid of the hook in the model. Also saves a view of the hook's position and quaternion from MjData, so that they can be more easily updated later.

update(self, pos, quat):

Inputs:

  • pos: an array of 3 floats, the position to be set
  • quat: an array of 4 floats, the quaternion to be set

Description:

It not only sets the position and orientation of the drone, but those of the hook too, so that the hook follows the drone. This will need to be changed later when there will be a sensor on the hook of the real drone, so that the hook angle can be changed too.

classes/trajectory_base.py

class TrajectoryBase

Description:

Base class for drone trajectories. Contains the evaluate method that trajectories must implement. This acts as an interface so that it'd be easy to equip drones with different trajectories. Each drone has a Trajectory member variable whose evaluate() method is called in the drone's update(), and the evaluate() always returns the required format for the drone's controllers.

__init__(self):

Description:

Initializes the dictionary that will be the output of each evaluate() call.

evaluate(self, state, i, time, control_step) -> dict:

Inputs:

  • state: state of the vehicle when trajectory is being evaluated
  • i: the loop variable
  • time: time that has passed since the simulation started
  • control_step: how frequently control of the vehicles is updated

Description:

The method, that derived classes must implement. It evaluates the trajectory and returns the input parameters (the dictionary initialized in the constructor) for the controllers.

classes/airflow_sampler.py

class AirflowSampler

Description:

Computes the force and torque on the payload due to the airflow generated by the propellers. The calculations are based on the air pressure and air velocity look-up tables that have been generated by CFD software. The sampler currently only supports box shaped payloads whose sides have been meshed (divided into small squares that'll be used to calculate the forces acting on them).

__init__(self, data_file_name_pressure : str, owning_drone : Drone, data_file_name_velocity: str = None):

Inputs:

  • data_file_name_pressure: path to the pressure look-up table
  • owning_drone: the drone to whose coordinate frame the sampler is attached
  • data_file_name_velocity: optional, path to the velocity look-up table... if not provided, only the pressure lut is used

Description:

Initializes an airflow sampler with the look-up tables provided as inputs. Assumes that in the look-up tables, each value is contained by a 1cm^3 small cube, and the overall cube size is determined this way.

Methods

get_transformed_vertices(self):

Description:

Returns the coordinates of the sampler cube in world coordinates.

get_payload_offset_z(self):

Description:

Returns the offset of the payload in centimeters. The payload is shifted with respect to the sampler cube because the cube is too small at the moment.

get_payload_offset_z_meter(self):

Description:

Returns the offset of the payload in meters.

generate_forces_opt(self, payload : Payload):

Inputs:

  • payload: the payload on which the airflow needs to be evaluated

Description:

Computes the forces and torques generated by the airflow on the payload with matrix multiplications. Runs faster than the original implementation by two orders of magnitude. Since the class only supports box shaped payloads for the moment, there's only one normal vector per side, because it'd be unnecessary to store the same normal for each tiny rectangle. Later, when arbitrary shapes will be supported, a new data structure will be needed where each divided surface area has the following properties: area, normal vector, position vector pointing from the center of mass (of the payload) to the center of the small area.

util/xml_generator.py

class SceneXmlGenerator

Description:

This class uses xml.etree.ElementTree to generate a MuJoCo model as an xml file.

__init__(self, base_scene_filename):

Inputs:

  • base_scene_filename: the name of the base xml file that contains default classes, assets (such as meshes, textures etc.) and some geoms whose position remains fixed throughout the simulation

Description:

Initializes ElementTree and the counters of the objects.

Methods

All the add methods follow the same pattern. If one of these methods is called, it adds one object to the xml model with varying degree of hard-coded information. The information that is not hard-coded in these methods, are the inputs (position, orientation, etc.) that can be different from object to object. Some objects can only be added to the model once like Sztaki, hospital, post office, parking lot, airport. And some objects can be added to the model multiple times, like drones, landing zones, poles, payloads.

save_xml(self, file_name):

Description:

Saves the xml model to hard disk at the specified file name.

Scripts

scripts/build_scene.py

Contains methods for building a scene with either GUI's or automatically from Motive stream. It uses a PassiveDisplay for visual representation, a SceneXmlGenerator to generate the model as xml, and three GUI's for adding buildings, drones, and payloads. At startup it loads the base xml model for pretty much all of the scenarios, that contains the 3D model of the 6th floor of Sztaki with the roundabout carpet. This xml gets included in all the generated xml's. For usage, see README.md.

Methods

add_building():

Description:

Creates and shows a BuildingInputGui pop-up window in which the user can select the kind of building, its position and orientation. It then adds the building to the xml model by the SceneXmlGenerator. Every time the user adds a new building, the xml is saved, and reloaded in MuJoCo, so that the user can see the new model visually. Some buildings can only be added once, like hospital and Sztaki, because there's only one of them in the real test environment too. Multiple poles and landing zones can be added however.

add_drone():

Description:

Creates and shows a DroneInputGui pop-up window in which the user can select the tye of drone, its position and orientation. Then, it adds the desired drone to the xml model via the SceneXmlGenerator, saves the xml, and gets the display to reload it so that the user can see the changes visually.

add_load():

Description:

Creates and shows a PayloadInputGui pop-up window in which the user can select the color of the payload, its mass, size, position and orientation. It is then added to the xml model via SceneXmlGenerator, the xml is saved to hard disk, and reloaded in PassiveDisplay to make the changes visible to the user.

clear_scene():

Description:

Reinitializes the model to the base xml model and clears counters.

build_from_optitrack():

Description:

Tries to build the entire scene from Motive stream. For this, all the buildings and vehichles that appear in the scene must be ticked in the assets pane of Motive. Objects that are a plane - like airport - have a fixed z coordinate of 0.01, so that they'd be visible. The z coordinate of the position from Motive is generally not used, except for the drones. A landing zone is put underneath each drone, because landing zones don't have markers.

main():

Description:

Binds the above mentioned methods to keyboard events of PassiveDisplay, and runs the PassiveDisplay instance.

scripts/load_and_display_scene.py

Description:

Extends PassiveDisplay with a method to load models. The script demonstrates how to use the PassiveDisplay class. Initially the base scene model is loaded.

Methods

load_model():

Description:

Creates and shows a pop-up file dialog in which the user can select any MuJoCo xml models. Once a model is selected, it is loaded into the PassiveDisplay instance.

main():

Description:

Binds the load_model() method to PassiveDisplay's l keyboard event and runs the PassiveDisplay instance.

Helper methods

util/mujoco_helper.py

get_joint_name_list(mjmodel: mujoco.MjModel):

Description:

Goes through a loaded MuJoCo model, and creates a list of valid joint names that appear in the model.

get_geom_name_list(mjmodel: mujoco.MjModel):

Description:

Goes through a loaded MuJoCo model, and creates a list of valid geom names that appear in the model.

get_body_name_list(mjmodel: mujoco.MjModel):

Description:

Goes through a loaded MuJoCo model, and creates a list of valid body names that appear in the model.

update_onboard_cam(drone_qpos, cam, azim_filter_sin=None, azim_filter_cos=None, elev_filter_sin=None, elev_filter_cos=None):

Inputs:

  • drone_qpos: a 7 element array or list, first 3 of which are position, remaining 4 are quaternion of a drone
  • cam: camera to be updated
  • azim_filter_sin: a LiveLFilter low-pass filter for the sine component of the camera azimuth angle
  • azim_filter_cos: a LiveLFilter low-pass filter for the cosine component of the camera azimuth angle
  • elev_filter_sin: a LiveLFilter low-pass filter for the sine component of the camera elevation angle
  • elev_filter_cos: a LiveLFilter low-pass filter for the cosine component of the camera elevation angle

Description:

Update the position and orientation of the camera that follows the drone from behind qpos is the array in which the position and orientation of all the drones are stored

Smoothing the 2 angle signals with 4 low-pass filters so that the camera would not shake. It's not enough to only filter the angle signal, because when the drone turns, the angle might jump from 180 degrees to -180 degrees, and the filter tries to smooth out the jump (the camera ends up turning a 360). Instead, take the sine and the cosine of the angle, filter them, and convert them back with atan2().

Clone this wiki locally