--- a
+++ b/SAC/IK_Framework_Mujoco.py
@@ -0,0 +1,373 @@
+#Minimal implementation of the RL_Framework_Mujoco environment for CMA-ES and IK Optimization
+#for finding the inital pose of the musculoskeletal model
+#and for trajectory visualization
+
+from gym import error, spaces
+from gym.utils import seeding
+import numpy as np
+from os import path
+import gym
+import pickle
+
+import numpy as np
+from gym import utils
+from . import sensory_feedback_specs
+from . import kinematics_preprocessing_specs
+
+import ipdb
+
+try:
+    import mujoco_py
+except ImportError as e:
+    raise error.DependencyNotInstalled("{}. (HINT: you need to install mujoco_py, and also perform the setup instructions here: https://github.com/openai/mujoco-py/.)".format(e))
+
+DEFAULT_SIZE = 500
+
+class MujocoEnv(gym.Env):
+    """Superclass for all MuJoCo environments.
+    """
+    def __init__(self, model_path, condition_to_sim, cond_tpoint, args):
+
+        self.model_path = model_path
+        self.initial_pose_path = args.initial_pose_path
+        self.kinematics_path = args.kinematics_path
+        self.cond_to_sim = condition_to_sim
+        self.cond_tpoint = cond_tpoint
+
+         #setup the model
+        self.model = mujoco_py.load_model_from_path(model_path)
+
+        #setup the simulator and data object
+        self.sim = mujoco_py.MjSim(self.model)
+        self.data = self.sim.data
+
+        #Save all the sensory feedback specs for use in the later functions
+        self.sfs_proprioceptive_feedback = args.proprioceptive_feedback
+        self.sfs_muscle_forces = args.muscle_forces
+        self.sfs_joint_feedback = args.joint_feedback
+        self.sfs_visual_feedback = args.visual_feedback
+        self.sfs_visual_feedback_bodies = args.visual_feedback_bodies
+        self.sfs_visual_distance_bodies = args.visual_distance_bodies
+        self.sfs_visual_velocity = args.visual_velocity
+        self.sfs_sensory_delay_timepoints = args.sensory_delay_timepoints
+
+        #Load the kin_train
+        # Load the experimental kinematics x and y coordinates from the data
+        with open(self.kinematics_path + '/kinematics.pkl', 'rb') as f:
+            kin_train_test = pickle.load(f)
+
+        kin_train = kin_train_test['train'] #[num_conds][num_targets, num_coords, timepoints]
+
+        #Randomly sample the number of targets from a condition = 0
+        num_targets = kin_train[0].shape[0]
+
+        #Find the qpos corresponding to the musculoskeletal model and targets
+        self.qpos_idx_musculo = np.array(list(range(0, self.model.nq)))
+        self.qpos_idx_targets = []
+        for musculo_targets in kinematics_preprocessing_specs.musculo_target_joints:
+        	joint_idx = self.model.get_joint_qpos_addr(musculo_targets)
+        	self.qpos_idx_targets.append(joint_idx)
+        	
+        #Delete the corresponding index from the qpos_idx_musculo
+        self.qpos_idx_musculo = np.delete(self.qpos_idx_musculo, self.qpos_idx_targets).tolist()
+
+        #Set the simulation timestep
+        if args.sim_dt != 0:
+            self.model.opt.timestep = args.sim_dt
+
+
+        self.n_fixedsteps = args.n_fixedsteps
+        self.radius = args.trajectory_scaling
+        self.center = args.center
+
+        #Kinematics preprocessing for training and testing kinematics
+        #Preprocess training kinematics
+        for i_target in range(kin_train[0].shape[0]):
+            for i_cond in range(len(kin_train)):
+                for i_coord in range(kin_train[i_cond].shape[1]):
+                    kin_train[i_cond][i_target, i_coord, :] = kin_train[i_cond][i_target, i_coord, :] / self.radius[i_target]
+                    kin_train[i_cond][i_target, i_coord, :] = kin_train[i_cond][i_target, i_coord, :] + self.center[i_target][i_coord]
+
+
+        self.kin_train = kin_train
+        self.kin_to_sim = self.kin_train
+
+        self.current_cond_to_sim = self.cond_to_sim
+
+        self.upd_theta(self.cond_tpoint)
+
+
+        self.viewer = None 
+        self._viewers = {}
+
+        self.metadata = {
+            'render.modes': ['human', 'rgb_array', 'depth_array'],
+            'video.frames_per_second': int(np.round(1.0 / self.dt))
+        }
+
+        #init_qpos and init_qvel do not contain target qpos
+        #and contain only the musculo qpos
+        #If the initial qpos and qvel are provided by the user
+        if path.isfile(self.initial_pose_path + '/qpos.npy'):
+            init_qpos = np.load(self.initial_pose_path + '/qpos.npy')
+            init_qvel = np.load(self.initial_pose_path + '/qvel.npy')
+        
+        #else use the default initial pose of xml model
+        else:
+            init_qpos = self.sim.data.qpos.flat.copy()[self.qpos_idx_musculo]
+            init_qvel = self.sim.data.qvel.flat.copy()[self.qpos_idx_musculo]
+
+        #Get the qpos of musclo + targets
+        musculo_qpos = self.sim.data.qpos.flat.copy() 
+        musculo_qvel = self.sim.data.qvel.flat.copy()
+
+        #Set the musculo part to the saved initial qpos and qvel
+        musculo_qpos[self.qpos_idx_musculo] = init_qpos
+
+        self.init_qpos = musculo_qpos
+
+        #Set the initial state to init_qpos
+        #Set the state to the initial pose
+        self.set_state(self.init_qpos)
+
+        if self.sfs_visual_feedback == True:
+
+            #Save the xpos of the musculo bodies for visual vels
+            if len(self.sfs_visual_velocity) != 0:
+                self.prev_body_xpos = []
+                for musculo_body in self.sfs_visual_velocity:
+                    body_xpos = self.sim.data.get_body_xpos(musculo_body)
+                    self.prev_body_xpos = [*self.prev_body_xpos, *body_xpos]
+
+    #Return the qpos of only the musculoskeletal bodies and not the targets
+    def get_musculo_state(self):
+    	qpos_all = self.sim.data.qpos.flat.copy()
+    	qpos_musculo = qpos_all[self.qpos_idx_musculo]
+
+    	return qpos_musculo
+   
+
+    def set_state(self, qpos):
+        
+        assert qpos.shape == (self.model.nq, )
+        
+        old_state= self.sim.get_state()
+
+        new_state= mujoco_py.MjSimState(old_state.time, qpos, qpos*0,
+                                        old_state.act, old_state.udd_state)
+
+        self.sim.set_state(new_state)
+        self.sim.forward()
+
+    def set_state_musculo(self, qpos):
+        
+        old_state= self.sim.get_state()
+        qpos_all = self.sim.data.qpos.flat.copy()
+        qpos_all[self.qpos_idx_musculo] = qpos
+
+        assert qpos_all.shape == (self.model.nq, )
+
+        new_state= mujoco_py.MjSimState(old_state.time, qpos_all, qpos_all*0,
+                                        old_state.act, old_state.udd_state)
+
+        self.sim.set_state(new_state)
+        self.sim.forward()
+
+    def get_obs_musculo_bodies(self):
+    	#Returns the current xyz coords of the musculo bodies to be tracked
+    	musculo_body_state = []
+    	for musculo_body in kinematics_preprocessing_specs.musculo_tracking:
+    		current_xyz_coord = self.sim.data.get_body_xpos(musculo_body[0]).flat.copy() 
+    		musculo_body_state.append(current_xyz_coord)
+
+    	return np.array(musculo_body_state)  #[n_musculo_bodies, 3]
+
+
+    def get_obs_targets(self):
+    	#Returns the current xyz coords of the targets to be tracked
+    	musculo_target_state = []
+    	for musculo_body in kinematics_preprocessing_specs.musculo_tracking:
+    		current_xyz_coord = self.sim.data.get_body_xpos(musculo_body[1]).flat.copy() 
+    		musculo_target_state.append(current_xyz_coord)
+
+    	return np.array(musculo_target_state)  #[n_musculo_targets, 3]
+
+    def set_cond_to_simulate(self, i_condition, cond_timepoint):
+
+        #update the state variables
+        self.cond_to_sim = i_condition
+        self.cond_tpoint = cond_timepoint
+
+        self.current_cond_to_sim = self.cond_to_sim
+
+        self.upd_theta(self.cond_tpoint)
+
+
+    @property
+    def dt(self):
+        return self.model.opt.timestep 
+
+    def render(self,
+               mode='human',
+               width=DEFAULT_SIZE,
+               height=DEFAULT_SIZE,
+               camera_id=0,
+               camera_name=None):
+        if mode == 'rgb_array' or mode == 'depth_array':
+            if camera_id is not None and camera_name is not None:
+                raise ValueError("Both `camera_id` and `camera_name` cannot be"
+                                 " specified at the same time.")
+
+            no_camera_specified = camera_name is None and camera_id is None
+            if no_camera_specified:
+                camera_name = 'track'
+
+            if camera_id is None and camera_name in self.model._camera_name2id:
+                camera_id = self.model.camera_name2id(camera_name)
+
+            self._get_viewer(mode).render(width, height, camera_id=camera_id)
+
+        if mode == 'rgb_array':
+            # window size used for old mujoco-py:
+            data = self._get_viewer(mode).read_pixels(width, height, depth=False)
+            # original image is upside-down, so flip it
+            return data[::-1, :, :]
+        elif mode == 'depth_array':
+            self._get_viewer(mode).render(width, height)
+            # window size used for old mujoco-py:
+            # Extract depth part of the read_pixels() tuple
+            data = self._get_viewer(mode).read_pixels(width, height, depth=True)[1]
+            # original image is upside-down, so flip it
+            return data[::-1, :]
+        elif mode == 'human':
+            self._get_viewer(mode).render()
+
+    def close(self):
+        if self.viewer is not None:
+            # self.viewer.finish()
+            self.viewer = None
+            self._viewers = {}
+
+    def _get_viewer(self, mode):
+        self.viewer = self._viewers.get(mode)
+        if self.viewer is None:
+            if mode == 'human':
+                self.viewer = mujoco_py.MjViewer(self.sim)
+            elif mode == 'rgb_array' or mode == 'depth_array':
+                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
+
+            self.viewer_setup()
+            self._viewers[mode] = self.viewer
+        return self.viewer
+
+    def get_body_com(self, body_name):
+        return self.data.get_body_xpos(body_name).copy()
+
+
+class Muscle_Env(MujocoEnv):
+
+    def __init__(self, model_path, condition_to_sim, cond_tpoint, args):
+        MujocoEnv.__init__(self, model_path, condition_to_sim, cond_tpoint, args)
+
+    def viewer_setup(self):
+        self.viewer.cam.trackbodyid = 0
+
+    def upd_theta(self, cond_timepoint):
+
+        coords_to_sim = self.kin_to_sim[self.current_cond_to_sim] #[num_targets, num_coords, timepoints]
+
+        assert cond_timepoint < coords_to_sim.shape[-1]
+        
+        crnt_state = self.sim.get_state()
+
+        for i_target in range(self.kin_to_sim[self.current_cond_to_sim].shape[0]):
+            if kinematics_preprocessing_specs.xyz_target[i_target][0]:
+                x_joint_idx= self.model.get_joint_qpos_addr(f"box:x{i_target}")
+                crnt_state.qpos[x_joint_idx] = coords_to_sim[i_target, 0, cond_timepoint]
+
+
+            if kinematics_preprocessing_specs.xyz_target[i_target][1]:
+                y_joint_idx= self.model.get_joint_qpos_addr(f"box:y{i_target}")
+                crnt_state.qpos[y_joint_idx] = coords_to_sim[i_target, kinematics_preprocessing_specs.xyz_target[i_target][0], cond_timepoint]
+
+            if kinematics_preprocessing_specs.xyz_target[i_target][2]:
+                z_joint_idx= self.model.get_joint_qpos_addr(f"box:z{i_target}")
+                crnt_state.qpos[z_joint_idx] = coords_to_sim[i_target, kinematics_preprocessing_specs.xyz_target[i_target][0] + kinematics_preprocessing_specs.xyz_target[i_target][1], cond_timepoint]
+
+
+        #Now set the state
+        self.set_state(crnt_state.qpos)
+
+
+    def _get_obs(self):
+        
+        sensory_feedback = []
+        if self.sfs_proprioceptive_feedback == True:
+            muscle_lens = self.sim.data.actuator_length.flat.copy()
+            muscle_vels = self.sim.data.actuator_velocity.flat.copy()
+
+            #process through the given function for muscle lens and muscle vels
+            muscle_lens, muscle_vels = sensory_feedback_specs.process_proprioceptive(muscle_lens, muscle_vels)
+            sensory_feedback = [*sensory_feedback, *muscle_lens, *muscle_vels]
+
+
+        if self.sfs_muscle_forces == True:
+            actuator_forces = self.sim.data.qfrc_actuator.flat.copy()
+
+            #process
+            actuator_forces = sensory_feedback_specs.process_muscle_forces(actuator_forces)
+            sensory_feedback = [*sensory_feedback, *actuator_forces]
+
+
+        if self.sfs_joint_feedback == True:
+            sensory_qpos = self.sim.data.qpos.flat.copy()
+            sensory_qvel = self.sim.data.qvel.flat.copy()
+
+            sensory_qpos, sensory_qvel = sensory_feedback_specs.process_joint_feedback(sensory_qpos, sensory_qvel)
+            sensory_feedback = [*sensory_feedback, *sensory_qpos, *sensory_qvel]
+
+
+        if self.sfs_visual_feedback == True:
+            
+            #Check if the user specified the musculo bodies to be included
+            assert len(self.sfs_visual_feedback_bodies) != 0
+
+            visual_xyz_coords = []
+            for musculo_body in self.sfs_visual_feedback_bodies:
+                visual_xyz_coords = [*visual_xyz_coords, *self.sim.data.get_body_xpos(musculo_body)]
+
+            visual_xyz_coords = sensory_feedback_specs.process_visual_position(visual_xyz_coords)
+            sensory_feedback = [*sensory_feedback, *visual_xyz_coords]
+
+        if len(self.sfs_visual_distance_bodies) != 0:
+            visual_xyz_distance = []
+            for musculo_tuple in self.sfs_visual_distance_bodies:
+                body0_xyz = self.sim.data.get_body_xpos(musculo_tuple[0])
+                body1_xyz = self.sim.data.get_body_xpos(musculo_tuple[1])
+                tuple_dist = np.abs(body0_xyz - body1_xyz).tolist()
+                visual_xyz_distance = [*visual_xyz_distance, *tuple_dist]
+
+            #process
+            visual_xyz_distance = sensory_feedback_specs.process_visual_distance(visual_xyz_distance)
+            sensory_feedback = [*sensory_feedback, *visual_xyz_distance]
+
+        #Save the xpos of the musculo bodies for visual vels
+        if len(self.sfs_visual_velocity) != 0:
+
+            #Find the visual vels after the simulation
+            current_body_xpos = []
+            for musculo_body in self.sfs_visual_velocity:
+                body_xpos = self.sim.data.get_body_xpos(musculo_body)
+                current_body_xpos = [*current_body_xpos, *body_xpos]
+
+            #Find the velocity
+            visual_vels = (np.abs(np.array(self.prev_body_xpos) - np.array(current_body_xpos)) / self.dt).tolist()
+            
+            #process visual velocity feedback
+            visual_vels = sensory_feedback_specs.process_visual_velocity(visual_vels)
+            sensory_feedback= [*sensory_feedback, *visual_vels]
+
+            #update the self.prev_body_xpos
+            self.prev_body_xpos = current_body_xpos
+
+        return np.array(sensory_feedback)
\ No newline at end of file