Files
rotor_py_control/rotorpy/environments.py
2024-08-06 13:21:48 -04:00

187 lines
9.9 KiB
Python

import numpy as np
from scipy.spatial.transform import Rotation
import matplotlib.pyplot as plt
import time as clk
from rotorpy.simulate import simulate
from rotorpy.utils.plotter import *
from rotorpy.world import World
from rotorpy.utils.postprocessing import unpack_sim_data
import os
class Environment():
"""
Sandbox represents an instance of the simulation environment containing a unique vehicle,
controller, trajectory generator, wind profile.
"""
def __init__(self, vehicle, # vehicle object, must be specified.
controller, # controller object, must be specified.
trajectory, # trajectory object, must be specified.
wind_profile = None, # wind profile object, if none is supplied it will choose no wind.
imu = None, # imu sensor object, if none is supplied it will choose a default IMU sensor.
mocap = None, # mocap sensor object, if none is supplied it will choose a default mocap.
world = None, # The world object
estimator = None, # estimator object
sim_rate = 100, # The update frequency of the simulator in Hz
safety_margin = 0.25, # The radius of the safety region around the robot.
):
self.sim_rate = sim_rate
self.vehicle = vehicle
self.controller = controller
self.trajectory = trajectory
self.safety_margin = safety_margin
if world is None:
# If no world is specified, assume that it means that the intended world is free space.
wbound = 3
self.world = World.empty((-wbound, wbound, -wbound,
wbound, -wbound, wbound))
else:
self.world = world
if wind_profile is None:
# If wind is not specified, default to no wind.
from rotorpy.wind.default_winds import NoWind
self.wind_profile = NoWind()
else:
self.wind_profile = wind_profile
if imu is None:
# In the event of specified IMU, default to 0 bias with white noise with default parameters as specified below.
from rotorpy.sensors.imu import Imu
self.imu = Imu(p_BS = np.zeros(3,),
R_BS = np.eye(3),
sampling_rate=sim_rate)
else:
self.imu = imu
if mocap is None:
# If no mocap is specified, set a default mocap.
# Default motion capture properties. Pretty much made up based on qualitative comparison with real data from Vicon.
mocap_params = {'pos_noise_density': 0.0005*np.ones((3,)), # noise density for position
'vel_noise_density': 0.0010*np.ones((3,)), # noise density for velocity
'att_noise_density': 0.0005*np.ones((3,)), # noise density for attitude
'rate_noise_density': 0.0005*np.ones((3,)), # noise density for body rates
'vel_artifact_max': 5, # maximum magnitude of the artifact in velocity (m/s)
'vel_artifact_prob': 0.001, # probability that an artifact will occur for a given velocity measurement
'rate_artifact_max': 1, # maximum magnitude of the artifact in body rates (rad/s)
'rate_artifact_prob': 0.0002 # probability that an artifact will occur for a given rate measurement
}
from rotorpy.sensors.external_mocap import MotionCapture
self.mocap = MotionCapture(sampling_rate=sim_rate, mocap_params=mocap_params, with_artifacts=False)
else:
self.mocap = mocap
if estimator is None:
# In the likely case where an estimator is not supplied, default to the null state estimator.
from rotorpy.estimators.nullestimator import NullEstimator
self.estimator = NullEstimator()
else:
self.estimator = estimator
return
def run(self, t_final = 10, # The maximum duration of the environment in seconds
use_mocap = False, # boolean determines if the controller should use
terminate = False,
plot = False, # Boolean: plots the vehicle states and commands
plot_mocap = True, # Boolean: plots the motion capture pose and twist measurements
plot_estimator = True, # Boolean: plots the estimator filter states and covariance diagonal elements
plot_imu = True, # Boolean: plots the IMU measurements
animate_bool = False, # Boolean: determines if the animation of vehicle state will play.
animate_wind = False, # Boolean: determines if the animation will include a wind vector.
verbose = False, # Boolean: will print statistics regarding the simulation.
fname = None # Filename is specified if you want to save the animation. Default location is the home directory.
):
"""
Run the simulator
"""
self.t_step = 1/self.sim_rate
self.t_final = t_final
self.t_final = t_final
self.terminate = terminate
self.use_mocap = use_mocap
start_time = clk.time()
(time, state, control, flat, imu_measurements, imu_gt, mocap_measurements, state_estimate, exit) = simulate(self.world,
self.vehicle.initial_state,
self.vehicle,
self.controller,
self.trajectory,
self.wind_profile,
self.imu,
self.mocap,
self.estimator,
self.t_final,
self.t_step,
self.safety_margin,
self.use_mocap,
terminate=self.terminate,
)
if verbose:
# Print relevant statistics or simulator status indicators here
print('-------------------RESULTS-----------------------')
print('SIM TIME -- %3.2f seconds | WALL TIME -- %3.2f seconds' % (min(self.t_final, time[-1]) , (clk.time()-start_time)))
print('EXIT STATUS -- '+exit.value)
self.result = dict(time=time, state=state, control=control, flat=flat, imu_measurements=imu_measurements, imu_gt=imu_gt, mocap_measurements=mocap_measurements, state_estimate=state_estimate, exit=exit)
visualizer = Plotter(self.result, self.world)
# Remove gif or mp4 in filename if it exists (the respective functions will add appropriate extensions)
if fname is not None:
if ".gif" in fname:
fname = fname.replace(".gif", "")
if ".mp4" in fname:
fname = fname.replace(".mp4", "")
if animate_bool:
# Do animation here
visualizer.animate_results(fname=fname, animate_wind=animate_wind)
if plot:
# Do plotting here
visualizer.plot_results(fname=fname,plot_mocap=plot_mocap,plot_estimator=plot_estimator,plot_imu=plot_imu)
if not animate_bool:
plt.show()
return self.result
def save_to_csv(self, fname):
"""
Save the simulation data in self.results to a file.
"""
if self.result is None:
print("Error: cannot save if no results have been generated! Aborting save.")
return
else:
if not ".csv" in fname:
fname = fname + ".csv"
path = os.path.join(os.path.dirname(__file__),'data_out',fname)
dataframe = unpack_sim_data(self.result)
dataframe.to_csv(path)
if __name__=="__main__":
from rotorpy.vehicles.crazyflie_params import quad_params
from rotorpy.trajectories.hover_traj import HoverTraj
from rotorpy.vehicles.multirotor import Multirotor
from rotorpy.controllers.quadrotor_control import SE3Control
sim = Environment(vehicle=Multirotor(quad_params),
controller=SE3Control(quad_params),
trajectory=HoverTraj(),
sim_rate=100
)
result = sim.run(t_final=1, plot=True)