from stable_baselines3 import PPO
from stable_baselines3.common.save_util import save_to_zip_file, load_from_zip_file
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
from my_environments import Ultrasound, FetchPush
from my_models.grippers import UltrasoundProbeGripper
from helper import register_gripper, plot_joint_pos, plot_forces_and_contact, plot_gripper_position
from demos import controller_demo, \
contact_btw_probe_and_body_demo, \
standard_mujoco_py_demo, \
drop_cube_on_body_demo, \
change_parameters_of_soft_body_demo, \
fetch_push_gym_demo
register_env(Ultrasound)
register_env(FetchPush)
register_gripper(UltrasoundProbeGripper)
## Simulations
#controller_demo('UR5e', save_data=False)
#contact_btw_probe_and_body_demo(1, 'main', save_data=False)
#standard_mujoco_py_demo()
#drop_cube_on_body_demo()
#change_parameters_of_soft_body_demo(3)
#fetch_push_gym_demo()
#plot_joint_pos('data/UR5e_joint_pos_controller_test.csv')
#plot_forces_and_contact('data/main_ee_forces_contact_btw_probe_and_body.csv', 'data/main_contact_contact_btw_probe_and_body.csv')
#plot_gripper_position('data/main_gripper_pos_contact_btw_probe_and_body.csv')
:param env_id: (str) the environment ID
:param options: (dict) additional arguments to pass to the specific environment class initializer
:param seed: (int) the inital seed for RNG
:param rank: (int) index of the subprocess
"""
def _init():
register_gripper(UltrasoundProbeGripper)
env = GymWrapper(suite.make(env_id, **options))
env = Monitor(env)
env.seed(seed + rank)
return env
set_random_seed(seed)
return _init
if __name__ == '__main__':
register_env(FetchPush)
# Environment specifications
env_id = 'FetchPush'
options = {}
options['robots'] = 'UR5e'
options['controller_configs'] = load_controller_config(default_controller='OSC_POSE')
options['gripper_types'] = 'UltrasoundProbeGripper'
options['has_renderer'] = False
options['has_offscreen_renderer'] = False
options['use_camera_obs'] = False
options['use_object_obs'] = True
options['control_freq'] = 50
options['render_camera'] = None
options['horizon'] = 1000
options['reward_shaping'] = True
:param seed: (int) the inital seed for RNG
:param rank: (int) index of the subprocess
"""
def _init():
env = gym.make(env_id, reward_type="dense")
env = gym.wrappers.FlattenObservation(env)
env = Monitor(env)
env.seed(seed + rank)
return env
set_random_seed(seed)
return _init
if __name__ == '__main__':
register_env(Ultrasound)
with open("rl_config.yaml", 'r') as stream:
config = yaml.safe_load(stream)
# Environment specifications
env_options = config["robosuite"]
env_id = env_options.pop("env_id")
# Settings for stable-baselines RL algorithm
sb_config = config["sb_config"]
training_timesteps = sb_config["total_timesteps"]
check_pt_interval = sb_config["check_pt_interval"]
num_cpu = sb_config["num_cpu"]
# Settings for stable-baselines policy