class RatBMI(BMILoop, LogExperiment):
status = dict(wait=dict(start_trial='feedback_on', stop=None),
feedback_on=dict(baseline_hit='periph_targets', stop=None),
periph_targets=dict(target_hit='check_reward',
timeout='noise_burst',
stop=None),
check_reward=dict(rewarded_target='reward',
unrewarded_target='feedback_pause'),
feedback_pause=dict(end_feedback_pause='wait'),
reward=dict(reward_end='wait'),
noise_burst=dict(noise_burst_end='noise_burst_timeout'),
noise_burst_timeout=dict(noise_burst_timeout_end='wait'))
#Flag for feedback on or not
feedback = False
prev_targ_hit = 't1'
timeout_time = traits.Float(30.)
noise_burst_time = traits.Float(3.)
noise_burst_timeout_time = traits.Float(1.)
reward_time = traits.Float(1., desc='reward time')
#Frequency range:
aud_freq_range = traits.Tuple((1000., 20000.))
plant_type = traits.OptionsList(*plantlist,
desc='',
bmi3d_input_options=plantlist.keys())
#Time to average over:
nsteps = traits.Float(10.)
feedback_pause = traits.Float(3.)
def __init__(self, *args, **kwargs):
super(RatBMI, self).__init__(*args, **kwargs)
if hasattr(self, 'decoder'):
print self.decoder
else:
self.decoder = kwargs['decoder']
dec_params = dict(nsteps=self.nsteps, freq_lim=self.aud_freq_range)
for k, (key, val) in enumerate(dec_params.items()):
print key, val, self.decoder.filt.dec_params[key]
assert self.decoder.filt.dec_params[key] == val
self.decoder.filt.init_from_task(self.decoder.n_units, **dec_params)
self.plant = plantlist[self.plant_type]
def init(self, *args, **kwargs):
self.add_dtype('cursor', 'f8', (2, ))
self.add_dtype('freq', 'f8', (2, ))
super(RatBMI, self).init()
self.decoder.count_max = self.feature_accumulator.count_max
def _cycle(self):
self.rat_cursor = self.decoder.filt.get_mean()
self.task_data['cursor'] = self.rat_cursor
self.task_data['freq'] = self.decoder.filt.F
self.decoder.cnt = self.feature_accumulator.count
self.decoder.feedback = self.feedback
super(RatBMI, self)._cycle()
# def move_plant(self):
# if self.feature_accumulator.count == self.feature_accumulator.count_max:
# print 'self.plant.drive from task.py'
# self.plant.drive(self.decoder)
def _start_wait(self):
return True
def _test_start_trial(self, ts):
return True
def _test_rewarded_target(self, ts):
if self.prev_targ_hit == 't1':
return False
elif self.prev_targ_hit == 't2':
return True
def _test_unrewarded_target(self, ts):
if self.prev_targ_hit == 't1':
return True
elif self.prev_targ_hit == 't2':
return False
def _start_feedback_pause(self):
self.feedback = False
def _test_end_feedback_pause(self, ts):
return ts > self.feedback_pause
def _start_reward(self):
print 'reward!'
def _start_feedback_on(self):
self.feedback = True
def _test_baseline_hit(self, ts):
if self.prev_targ_hit == 't1':
#Must go below baseline:
return self.rat_cursor <= self.decoder.filt.mid
elif self.prev_targ_hit == 't2':
#Must rise above baseline:
return self.rat_cursor >= self.decoder.filt.mid
else:
return False
def _test_target_hit(self, ts):
if self.rat_cursor >= self.decoder.filt.t1:
self.prev_targ_hit = 't1'
self.feedback = False
return True
elif self.rat_cursor <= self.decoder.filt.t2:
self.prev_targ_hit = 't2'
self.feedback = False
return True
else:
return False
def _test_timeout(self, ts):
return ts > self.timeout_time
def _test_noise_burst_end(self, ts):
return ts > self.noise_burst_time
def _test_noise_burst_timeout_end(self, ts):
return ts > self.noise_burst_timeout_time
def _start_noise_burst(self):
self.feedback = False
self.plant.play_white_noise()
def move_plant(self):
super(RatBMI, self).move_plant()
def get_current_assist_level(self):
return 0.
class ArmPlant(Window):
'''
This task creates a RobotArm object and allows it to move around the screen based on either joint or endpoint
positions. There is a spherical cursor at the end of the arm. The links of the arm can be visible or hidden.
'''
background = (0, 0, 0, 1)
arm_visible = traits.Bool(
True,
desc='Specifies whether entire arm is displayed or just endpoint')
cursor_radius = traits.Float(.5, desc="Radius of cursor")
cursor_color = (.5, 0, .5, 1)
arm_class = traits.OptionsList(*plantlist,
bmi3d_input_options=plantlist.keys())
starting_pos = (5, 0, 5)
def __init__(self, *args, **kwargs):
super(ArmPlant, self).__init__(*args, **kwargs)
self.cursor_visible = True
# Initialize the arm
self.arm = ik.test_3d
self.arm_vis_prev = True
if self.arm_class == 'CursorPlant':
pass
else:
self.dtype.append(('joint_angles', 'f8', (self.arm.num_joints, )))
self.dtype.append(('arm_visible', 'f8', (1, )))
self.add_model(self.arm)
## Declare cursor
self.dtype.append(('cursor', 'f8', (3, )))
self.cursor = Sphere(radius=self.cursor_radius,
color=self.cursor_color)
self.add_model(self.cursor)
self.cursor.translate(*self.arm.get_endpoint_pos(), reset=True)
def _cycle(self):
'''
Calls any update functions necessary and redraws screen. Runs 60x per second by default.
'''
## Run graphics commands to show/hide the arm if the visibility has changed
if self.arm_class != 'CursorPlant':
if self.arm_visible != self.arm_vis_prev:
self.arm_vis_prev = self.arm_visible
self.show_object(self.arm, show=self.arm_visible)
self.move_arm()
self.update_cursor()
if self.cursor_visible:
self.task_data['cursor'] = self.cursor.xfm.move.copy()
else:
#if the cursor is not visible, write NaNs into cursor location saved in file
self.task_data['cursor'] = np.array([np.nan, np.nan, np.nan])
if self.arm_class != 'CursorPlant':
if self.arm_visible:
self.task_data['arm_visible'] = 1
else:
self.task_data['arm_visible'] = 0
super(ArmPlant, self)._cycle()
## Functions to move the cursor using keyboard/mouse input
def get_mouse_events(self):
import pygame
events = []
for btn in pygame.event.get(
(pygame.MOUSEBUTTONDOWN, pygame.MOUSEBUTTONUP)):
events = events + [btn.button]
return events
def get_key_events(self):
import pygame
return pygame.key.get_pressed()
def move_arm(self):
'''
allows use of keyboard keys to test movement of arm. Use QW/OP for joint movements, arrow keys for endpoint movements
'''
import pygame
keys = self.get_key_events()
joint_speed = (np.pi / 6) / 60
hand_speed = .2
x, y, z = self.arm.get_endpoint_pos()
if keys[pygame.K_RIGHT]:
x = x - hand_speed
self.arm.set_endpoint_pos(np.array([x, 0, z]))
if keys[pygame.K_LEFT]:
x = x + hand_speed
self.arm.set_endpoint_pos(np.array([x, 0, z]))
if keys[pygame.K_DOWN]:
z = z - hand_speed
self.arm.set_endpoint_pos(np.array([x, 0, z]))
if keys[pygame.K_UP]:
z = z + hand_speed
self.arm.set_endpoint_pos(np.array([x, 0, z]))
if self.arm.num_joints == 2:
xz, xy = self.get_arm_joints()
e = np.array([xz[0], xy[0]])
s = np.array([xz[1], xy[1]])
if keys[pygame.K_q]:
s = s - joint_speed
self.set_arm_joints([e[0], s[0]], [e[1], s[1]])
if keys[pygame.K_w]:
s = s + joint_speed
self.set_arm_joints([e[0], s[0]], [e[1], s[1]])
if keys[pygame.K_o]:
e = e - joint_speed
self.set_arm_joints([e[0], s[0]], [e[1], s[1]])
if keys[pygame.K_p]:
e = e + joint_speed
self.set_arm_joints([e[0], s[0]], [e[1], s[1]])
if self.arm.num_joints == 4:
jts = self.get_arm_joints()
keyspressed = [
keys[pygame.K_q], keys[pygame.K_w], keys[pygame.K_e],
keys[pygame.K_r]
]
for i in range(self.arm.num_joints):
if keyspressed[i]:
jts[i] = jts[i] + joint_speed
self.set_arm_joints(jts)
def get_cursor_location(self):
return self.arm.get_endpoint_pos()
def set_arm_endpoint(self, pt, **kwargs):
self.arm.set_endpoint_pos(pt, **kwargs)
def set_arm_joints(self, angle_xz, angle_xy):
self.arm.set_intrinsic_coordinates(angle_xz, angle_xy)
def get_arm_joints(self):
return self.arm.get_intrinsic_coordinates()
def update_cursor(self):
'''
Update the cursor's location and visibility status.
'''
pt = self.get_cursor_location()
if pt is not None:
self.move_cursor(pt)
def move_cursor(self, pt):
''' Move the cursor object to the specified 3D location. '''
if not hasattr(self.arm, 'endpt_cursor'):
self.cursor.translate(*pt[:3], reset=True)
class LFP_Mod(BMILoop, Sequence, Window):
background = (0,0,0,1)
plant_visible = traits.Bool(True, desc='Specifies whether entire plant is displayed or just endpoint')
lfp_cursor_rad = traits.Float(.5, desc="length of LFP cursor")
lfp_cursor_color = (.5,0,.5,.75)
lfp_plant_type_options = plantlist.keys()
lfp_plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys())
window_size = traits.Tuple((1920*2, 1080), desc='window size')
lfp_frac_lims = traits.Tuple((0., 0.35), desc='fraction limits')
xlfp_frac_lims = traits.Tuple((-.7, 1.7), desc = 'x dir fraction limits')
lfp_control_band = traits.Tuple((25, 40), desc='beta power band limits')
lfp_totalpw_band = traits.Tuple((1, 100), desc='total power band limits')
xlfp_control_band = traits.Tuple((0, 5), desc = 'x direction band limits')
n_steps = traits.Int(2, desc='moving average for decoder')
powercap = traits.Float(1, desc="Timeout for total power above this")
zboundaries=(-12,12)
status = dict(
wait = dict(start_trial="lfp_target", stop=None),
lfp_target = dict(enter_lfp_target="lfp_hold", powercap_penalty="powercap_penalty", stop=None),
lfp_hold = dict(leave_early="lfp_target", lfp_hold_complete="reward", powercap_penalty="powercap_penalty"),
powercap_penalty = dict(powercap_penalty_end="lfp_target"),
reward = dict(reward_end="wait")
)
static_states = [] # states in which the decoder is not run
trial_end_states = ['reward']
lfp_cursor_on = ['lfp_target', 'lfp_hold']
#initial state
state = "wait"
#create settable traits
reward_time = traits.Float(.5, desc="Length of juice reward")
lfp_target_rad = traits.Float(3.6, desc="Length of targets in cm")
lfp_hold_time = traits.Float(.2, desc="Length of hold required at lfp targets")
lfp_hold_var = traits.Float(.05, desc="Length of hold variance required at lfp targets")
hold_penalty_time = traits.Float(1, desc="Length of penalty time for target hold error")
powercap_penalty_time = traits.Float(1, desc="Length of penalty time for timeout error")
# max_attempts = traits.Int(10, desc='The number of attempts at a target before\
# skipping to the next one')
session_length = traits.Float(0, desc="Time until task automatically stops. Length of 0 means no auto stop.")
#plant_hide_rate = traits.Float(0.0, desc='If the plant is visible, specifies a percentage of trials where it will be hidden')
lfp_target_color = (123/256.,22/256.,201/256.,.5)
mc_target_color = (1,0,0,.5)
target_index = -1 # Helper variable to keep track of which target to display within a trial
#tries = 0 # Helper variable to keep track of the number of failed attempts at a given trial.
cursor_visible = False # Determines when to hide the cursor.
no_data_count = 0 # Counter for number of missing data frames in a row
sequence_generators = ['lfp_mod_4targ']
def __init__(self, *args, **kwargs):
super(LFP_Mod, self).__init__(*args, **kwargs)
self.cursor_visible = True
print 'INIT FRAC LIMS: ', self.lfp_frac_lims
dec_params = dict(lfp_frac_lims = self.lfp_frac_lims,
xlfp_frac_lims = self.xlfp_frac_lims,
powercap = self.powercap,
zboundaries = self.zboundaries,
lfp_control_band = self.lfp_control_band,
lfp_totalpw_band = self.lfp_totalpw_band,
xlfp_control_band = self.xlfp_control_band,
n_steps = self.n_steps)
self.decoder.filt.init_from_task(**dec_params)
self.decoder.init_from_task(**dec_params)
self.lfp_plant = plantlist[self.lfp_plant_type]
if self.lfp_plant_type == 'inv_cursor_onedimLFP':
print 'MAKE SURE INVERSE GENERATOR IS ON'
self.plant_vis_prev = True
self.current_assist_level = 0
self.learn_flag = False
if hasattr(self.lfp_plant, 'graphics_models'):
for model in self.lfp_plant.graphics_models:
self.add_model(model)
# Instantiate the targets
'''
height and width on kinarm machine are 2.4. Here we make it 2.4/8*12 = 3.6
'''
lfp_target = VirtualSquareTarget(target_radius=self.lfp_target_rad, target_color=self.lfp_target_color)
self.targets = [lfp_target]
# Initialize target location variable
self.target_location_lfp = np.array([-100, -100, -100])
# Declare any plant attributes which must be saved to the HDF file at the _cycle rate
for attr in self.lfp_plant.hdf_attrs:
self.add_dtype(*attr)
def init(self):
self.plant = DummyPlant()
self.add_dtype('lfp_target', 'f8', (3,))
self.add_dtype('target_index', 'i', (1,))
self.add_dtype('powercap_flag', 'i',(1,))
for target in self.targets:
for model in target.graphics_models:
self.add_model(model)
super(LFP_Mod, self).init()
def _cycle(self):
'''
Calls any update functions necessary and redraws screen. Runs 60x per second.
'''
self.task_data['loop_time'] = self.iter_time()
self.task_data['lfp_target'] = self.target_location_lfp.copy()
self.task_data['target_index'] = self.target_index
#self.task_data['internal_decoder_state'] = self.decoder.filt.current_lfp_pos
self.task_data['powercap_flag'] = self.decoder.filt.current_powercap_flag
self.move_plant()
## Save plant status to HDF file, ###ADD BACK
lfp_plant_data = self.lfp_plant.get_data_to_save()
for key in lfp_plant_data:
self.task_data[key] = lfp_plant_data[key]
super(LFP_Mod, self)._cycle()
def move_plant(self):
feature_data = self.get_features()
# Save the "neural features" (e.g. spike counts vector) to HDF file
for key, val in feature_data.items():
self.task_data[key] = val
Bu = None
assist_weight = 0
target_state = np.zeros([self.decoder.n_states, self.decoder.n_subbins])
## Run the decoder
if self.state not in self.static_states:
neural_features = feature_data[self.extractor.feature_type]
self.call_decoder(neural_features, target_state, Bu=Bu, assist_level=assist_weight, feature_type=self.extractor.feature_type)
## Drive the plant to the decoded state, if permitted by the constraints of the plant
self.lfp_plant.drive(self.decoder)
self.task_data['decoder_state'] = decoder_state = self.decoder.get_state(shape=(-1,1))
return decoder_state
def run(self):
'''
See experiment.Experiment.run for documentation.
'''
# Fire up the plant. For virtual/simulation plants, this does little/nothing.
self.lfp_plant.start()
try:
super(LFP_Mod, self).run()
finally:
self.lfp_plant.stop()
##### HELPER AND UPDATE FUNCTIONS ####
def update_cursor_visibility(self):
''' Update cursor visible flag to hide cursor if there has been no good data for more than 3 frames in a row'''
prev = self.cursor_visible
if self.no_data_count < 3:
self.cursor_visible = True
if prev != self.cursor_visible:
self.show_object(self.cursor, show=True)
else:
self.cursor_visible = False
if prev != self.cursor_visible:
self.show_object(self.cursor, show=False)
def update_report_stats(self):
'''
see experiment.Experiment.update_report_stats for docs
'''
super(LFP_Mod, self).update_report_stats()
self.reportstats['Trial #'] = self.calc_trial_num()
self.reportstats['Reward/min'] = np.round(self.calc_events_per_min('reward', 120), decimals=2)
#### TEST FUNCTIONS ####
def _test_powercap_penalty(self, ts):
if self.decoder.filt.current_powercap_flag:
#Turn off power cap flag:
self.decoder.filt.current_powercap_flag = 0
return True
else:
return False
def _test_enter_lfp_target(self, ts):
'''
return true if the distance between center of cursor and target is smaller than the cursor radius in the x and z axis only
'''
cursor_pos = self.lfp_plant.get_endpoint_pos()
dx = np.linalg.norm(cursor_pos[0] - self.target_location_lfp[0])
dz = np.linalg.norm(cursor_pos[2] - self.target_location_lfp[2])
in_targ = False
if dx<= (self.lfp_target_rad/2.) and dz<= (self.lfp_target_rad/2.):
in_targ = True
return in_targ
# #return d <= (self.lfp_target_rad - self.lfp_cursor_rad)
# #If center of cursor enters target at all:
# return d <= (self.lfp_target_rad/2.)
# #New version:
# cursor_pos = self.lfp_plant.get_endpoint_pos()
# d = np.linalg.norm(cursor_pos[2] - self.target_location_lfp[2])
# d <= (self.lfp_target_rad - self.lfp_cursor_rad)
def _test_leave_early(self, ts):
'''
return true if cursor moves outside the exit radius
'''
cursor_pos = self.lfp_plant.get_endpoint_pos()
dx = np.linalg.norm(cursor_pos[0] - self.target_location_lfp[0])
dz = np.linalg.norm(cursor_pos[2] - self.target_location_lfp[2])
out_of_targ = False
if dx > (self.lfp_target_rad/2.) or dz > (self.lfp_target_rad/2.):
out_of_targ = True
#rad = self.lfp_target_rad - self.lfp_cursor_rad
#return d > rad
return out_of_targ
def _test_lfp_hold_complete(self, ts):
return ts>=self.lfp_hold_time_plus_var
# def _test_lfp_timeout(self, ts):
# return ts>self.timeout_time
def _test_powercap_penalty_end(self, ts):
if ts>self.powercap_penalty_time:
self.lfp_plant.turn_on()
return ts>self.powercap_penalty_time
def _test_reward_end(self, ts):
return ts>self.reward_time
def _test_stop(self, ts):
if self.session_length > 0 and (self.get_time() - self.task_start_time) > self.session_length:
self.end_task()
return self.stop
#### STATE FUNCTIONS ####
def _parse_next_trial(self):
self.targs = self.next_trial
def _start_wait(self):
super(LFP_Mod, self)._start_wait()
self.tries = 0
self.target_index = -1
#hide targets
for target in self.targets:
target.hide()
#get target locations for this trial
self._parse_next_trial()
self.chain_length = 1
self.lfp_hold_time_plus_var = self.lfp_hold_time + np.random.uniform(low=-1,high=1)*self.lfp_hold_var
def _start_lfp_target(self):
self.target_index += 1
self.target_index = 0
#only 1 target:
target = self.targets[0]
self.target_location_lfp = self.targs #Just one target.
target.move_to_position(self.target_location_lfp)
target.cue_trial_start()
def _start_lfp_hold(self):
#make next target visible unless this is the final target in the trial
idx = (self.target_index + 1)
if idx < self.chain_length:
target = self.targets[idx % 2]
target.move_to_position(self.targs[idx])
def _end_lfp_hold(self):
# change current target color to green
self.targets[self.target_index % 2].cue_trial_end_success()
def _start_timeout_penalty(self):
#hide targets
for target in self.targets:
target.hide()
self.tries += 1
self.target_index = -1
def _start_reward(self):
super(LFP_Mod, self)._start_reward()
#self.targets[self.target_index % 2].show()
def _start_powercap_penalty(self):
for target in self.targets:
target.hide()
self.lfp_plant.turn_off()
@staticmethod
def lfp_mod_4targ(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8):
'''Mimics beta modulation task from Kinarm Rig:
In Kinarm rig, the following linear transformations happen:
1. LFP cursor is calculated
2. mapped from fraction limits [0, .35] to [-1, 1] (unit_coordinates)
3. udp sent to kinarm machine and multiplied by 8
4. translated upward in the Y direction by + 2.5
This means, our targets which are at -8, [-0.75, 2.5, 5.75, 9.0]
must be translated down by 2.5 to: -8, [-3.25, 0. , 3.25, 6.5]
then divided by 8: -1, [-0.40625, 0. , 0.40625, 0.8125 ] in unit_coordinates
The radius is 1.2, which is 0.15 in unit_coordinates
Now, we map this to a new system:
- new_zero: (y1+y2) / 2
- new_scale: (y2 - y1) / 2
(([-0.40625, 0. , 0.40625, 0.8125 ]) * new_scale ) + new_zero
new_zero = 0
new_scale = 12
12 * [-0.40625, 0. , 0.40625, 0.8125 ]
= array([-4.875, 0. , 4.875, 9.75 ])
'''
new_zero = (boundaries[3]+boundaries[2]) / 2.
new_scale = (boundaries[3] - boundaries[2]) / 2.
kin_targs = np.array([-0.40625, 0. , 0.40625, 0.8125 ])
lfp_targ_y = (new_scale*kin_targs) + new_zero
for i in range(nblocks):
temp = lfp_targ_y.copy()
np.random.shuffle(temp)
if i==0:
z = temp.copy()
else:
z = np.hstack((z, temp))
#Fixed X axis:
x = np.tile(xaxis,(nblocks*4))
y = np.zeros(nblocks*4)
pairs = np.vstack([x, y, z]).T
return pairs
class LFP_Mod_plus_MC_reach(LFP_Mod_plus_MC_hold):
mc_cursor_radius = traits.Float(.5, desc="Radius of cursor")
mc_target_radius = traits.Float(3, desc="Radius of MC target")
mc_cursor_color = (.5,0,.5,1)
mc_plant_type_options = plantlist.keys()
mc_plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys())
origin_hold_time = traits.Float(.2, desc="Hold time in center")
mc_periph_holdtime = traits.Float(.2, desc="Hold time in center")
mc_timeout_time = traits.Float(10, desc="Time allowed to go between targets")
exclude_parent_traits = ['goal_cache_block'] #redefine this to NOT include marker_num, marker_count
marker_num = traits.Int(14,desc='Index')
marker_count = traits.Int(16,desc='Num of markers')
scale_factor = 3.0 #scale factor for converting hand movement to screen movement (1cm hand movement = 3.5cm cursor movement)
wait_flag = 1
# NOTE!!! The marker on the hand was changed from #0 to #14 on
# 5/19/13 after LED #0 broke. All data files saved before this date
# have LED #0 controlling the cursor.
limit2d = 1
# state_file = open("/home/helene/preeya/tot_pw.txt","w")
state_cnt = 0
status = dict(
wait = dict(start_trial="origin", stop=None),
origin = dict(enter_origin="origin_hold", stop=None),
origin_hold = dict(origin_hold_complete="lfp_target",leave_origin="hold_penalty", stop=None),
lfp_target = dict(enter_lfp_target="lfp_hold", leave_origin="hold_penalty", powercap_penalty="powercap_penalty", stop=None),
lfp_hold = dict(leave_early="lfp_target", lfp_hold_complete="mc_target", leave_origin="hold_penalty",powercap_penalty="powercap_penalty"),
mc_target = dict(enter_mc_target='mc_hold',mc_timeout="timeout_penalty", stop=None),
mc_hold = dict(leave_periph_early='hold_penalty',mc_hold_complete="reward"),
powercap_penalty = dict(powercap_penalty_end="origin"),
timeout_penalty = dict(timeout_penalty_end="wait"),
hold_penalty = dict(hold_penalty_end="origin"),
reward = dict(reward_end="wait"),
)
static_states = ['origin'] # states in which the decoder is not run
trial_end_states = ['reward', 'timeout_penalty']
lfp_cursor_on = ['lfp_target', 'lfp_hold', 'reward']
sequence_generators = ['lfp_mod_plus_MC_reach', 'lfp_mod_plus_MC_reach_INV']
def __init__(self, *args, **kwargs):
# import pickle
# decoder = pickle.load(open('/storage/decoders/cart20141216_03_cart_new2015_2.pkl'))
# self.decoder = decoder
super(LFP_Mod_plus_MC_reach, self).__init__(*args, **kwargs)
mc_origin = VirtualCircularTarget(target_radius=self.mc_target_radius, target_color=RED)
mc_periph = VirtualCircularTarget(target_radius=self.mc_target_radius, target_color=RED)
lfp_target = VirtualSquareTarget(target_radius=self.lfp_target_rad, target_color=self.lfp_target_color)
self.targets = [lfp_target, mc_origin, mc_periph]
# #Should be unnecessary:
# for target in self.targets:
# for model in target.graphics_models:
# self.add_model(model)
# self.lfp_plant = plantlist[self.lfp_plant_type]
# if hasattr(self.lfp_plant, 'graphics_models'):
# for model in self.lfp_plant.graphics_models:
# self.add_model(model)
# self.mc_plant = plantlist[self.mc_plant_type]
# if hasattr(self.mc_plant, 'graphics_models'):
# for model in self.mc_plant.graphics_models:
# self.add_model(model)
def _parse_next_trial(self):
t = self.next_trial
self.lfp_targ = t['lfp']
self.mc_targ_orig = t['origin']
self.mc_targ_periph = t['periph']
def _start_mc_target(self):
#Turn off LFP things
self.lfp_plant.turn_off()
self.targets[0].hide()
self.targets[1].hide()
target = self.targets[2] #MC target
self.target_location_mc = self.mc_targ_periph
target.move_to_position(self.target_location_mc)
target.cue_trial_start()
def _test_enter_mc_target(self,ts):
cursor_pos = self.mc_plant.get_endpoint_pos()
d = np.linalg.norm(cursor_pos - self.target_location_mc)
return d <= (self.mc_target_radius - self.mc_cursor_radius)
def _test_mc_timeout(self, ts):
return ts>self.mc_timeout_time
def _test_leave_periph_early(self, ts):
cursor_pos = self.mc_plant.get_endpoint_pos()
d = np.linalg.norm(cursor_pos - self.target_location_mc)
rad = self.mc_target_radius - self.mc_cursor_radius
return d > rad
def _test_mc_hold_complete(self, ts):
return ts>self.mc_periph_holdtime
def _timeout_penalty_end(self, ts):
print 'timeout', ts
#return ts > 1.
return True
def _end_mc_hold(self):
self.targets[2].cue_trial_end_success()
# def _cycle(self):
# if self.state_cnt < 3600*3:
# self.state_cnt +=1
# s = "%s\n" % self.state
# self.state_file.write(str(s))
# if self.state_cnt == 3600*3:
# self.state_file.close()
# super(LFP_Mod_plus_MC_reach, self)._cycle()
def _start_reward(self):
super(LFP_Mod_plus_MC_reach, self)._start_reward()
lfp_targ = self.targets[0]
mc_orig = self.targets[1]
lfp_targ.hide()
mc_orig.hide()
@staticmethod
def lfp_mod_plus_MC_reach(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8, target_distance=6, n_mc_targets=4, mc_target_angle_offset=0,**kwargs):
new_zero = (boundaries[3]+boundaries[2]) / 2.
new_scale = (boundaries[3] - boundaries[2]) / 2.
kin_targs = np.array([-0.40625, 0. , 0.40625, 0.8125 ])
lfp_targ_y = (new_scale*kin_targs) + new_zero
for i in range(nblocks):
temp = lfp_targ_y.copy()
np.random.shuffle(temp)
if i==0:
z = temp.copy()
else:
z = np.hstack((z, temp))
#Fixed X axis:
x = np.tile(xaxis,(nblocks*4))
y = np.zeros(nblocks*4)
lfp = np.vstack([x, y, z]).T
origin = np.zeros(( lfp.shape ))
theta = []
for i in range(nblocks*4):
temp = np.arange(0, 2*np.pi, 2*np.pi/float(n_mc_targets))
np.random.shuffle(temp)
theta = theta + [temp]
theta = np.hstack(theta)
theta = theta + (mc_target_angle_offset*(np.pi/180.))
x = target_distance*np.cos(theta)
y = np.zeros(len(theta))
z = target_distance*np.sin(theta)
periph = np.vstack([x, y, z]).T
it = iter([dict(lfp=lfp[i,:], origin=origin[i,:], periph=periph[i,:]) for i in range(lfp.shape[0])])
if ('return_arrays' in kwargs.keys()) and kwargs['return_arrays']==True:
return lfp, origin, periph
else:
return it
@staticmethod
def lfp_mod_plus_MC_reach_INV(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8, target_distance=6, n_mc_targets=4, mc_target_angle_offset=0):
kw = dict(return_arrays=True)
lfp, origin, periph = LFP_Mod_plus_MC_reach.lfp_mod_plus_MC_reach(nblocks=nblocks, boundaries=boundaries, xaxis=xaxis, target_distance=target_distance,
n_mc_targets=n_mc_targets, mc_target_angle_offset=mc_target_angle_offset,**kw)
#Invert LFP:
lfp[:,2] = -1.0*lfp[:,2]
it = iter([dict(lfp=lfp[i,:], origin=origin[i,:], periph=periph[i,:]) for i in range(lfp.shape[0])])
return it
class LFP_Mod_plus_MC_hold(LFP_Mod):
mc_cursor_radius = traits.Float(.5, desc="Radius of cursor")
mc_target_radius = traits.Float(3, desc="Radius of MC target")
mc_cursor_color = (.5,0,.5,1)
mc_plant_type_options = plantlist.keys()
mc_plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys())
origin_hold_time = traits.Float(.2, desc="Hold time in center")
exclude_parent_traits = ['goal_cache_block'] #redefine this to NOT include marker_num, marker_count
marker_num = traits.Int(14,desc='Index')
marker_count = traits.Int(16,desc='Num of markers')
joystick_method = traits.Float(1,desc="1: Normal velocity, 0: Position control")
joystick_speed = traits.Float(20, desc="Radius of cursor")
move_while_in_center = traits.Float(1, desc="1 = update plant while in lfp_target, lfp_hold, 0 = don't update in these states")
scale_factor = 3.0 #scale factor for converting hand movement to screen movement (1cm hand movement = 3.5cm cursor movement)
wait_flag = 1
# NOTE!!! The marker on the hand was changed from #0 to #14 on
# 5/19/13 after LED #0 broke. All data files saved before this date
# have LED #0 controlling the cursor.
limit2d = 1
status = dict(
wait = dict(start_trial="origin", stop=None),
origin = dict(enter_origin="origin_hold", stop=None),
origin_hold = dict(origin_hold_complete="lfp_target",leave_origin="hold_penalty", stop=None),
lfp_target = dict(enter_lfp_target="lfp_hold", leave_origin="hold_penalty", powercap_penalty="powercap_penalty", stop=None),
lfp_hold = dict(leave_early="lfp_target", lfp_hold_complete="reward", leave_origin="hold_penalty", powercap_penalty="powercap_penalty",stop=None),
powercap_penalty = dict(powercap_penalty_end="origin"),
hold_penalty = dict(hold_penalty_end="origin",stop=None),
reward = dict(reward_end="wait")
)
static_states = ['origin'] # states in which the decoder is not run
trial_end_states = ['reward']
lfp_cursor_on = ['lfp_target', 'lfp_hold', 'reward']
sequence_generators = ['lfp_mod_4targ_plus_mc_orig']
def __init__(self, *args, **kwargs):
super(LFP_Mod_plus_MC_hold, self).__init__(*args, **kwargs)
if self.move_while_in_center>0:
self.no_plant_update_states = []
else:
self.no_plant_update_states = ['lfp_target', 'lfp_hold']
mc_origin = VirtualCircularTarget(target_radius=self.mc_target_radius, target_color=RED)
lfp_target = VirtualSquareTarget(target_radius=self.lfp_target_rad, target_color=self.lfp_target_color)
self.targets = [lfp_target, mc_origin]
self.mc_plant = plantlist[self.mc_plant_type]
if hasattr(self.mc_plant, 'graphics_models'):
for model in self.mc_plant.graphics_models:
self.add_model(model)
# Declare any plant attributes which must be saved to the HDF file at the _cycle rate
for attr in self.mc_plant.hdf_attrs:
self.add_dtype(*attr)
self.target_location_mc = np.array([-100, -100, -100])
self.manual_control_type = None
self.current_pt=np.zeros([3]) #keep track of current pt
self.last_pt=np.zeros([3])
def init(self):
self.add_dtype('mc_targ', 'f8', (3,)) ###ADD BACK
super(LFP_Mod_plus_MC_hold, self).init()
def _cycle(self):
'''
Calls any update functions necessary and redraws screen. Runs 60x per second.
'''
self.task_data['mc_targ'] = self.target_location_mc.copy()
mc_plant_data = self.mc_plant.get_data_to_save()
for key in mc_plant_data:
self.task_data[key] = mc_plant_data[key]
super(LFP_Mod_plus_MC_hold, self)._cycle()
def _parse_next_trial(self):
t = self.next_trial
self.lfp_targ = t['lfp']
self.mc_targ_orig = t['origin']
def _start_origin(self):
if self.wait_flag:
self.origin_hold_time_store = self.origin_hold_time
self.origin_hold_time = 3
self.wait_flag = 0
else:
self.origin_hold_time = self.origin_hold_time_store
#only 1 target:
target = self.targets[1] #Origin
self.target_location_mc = self.mc_targ_orig #Origin
target.move_to_position(self.target_location_mc)
target.cue_trial_start()
#Turn off lfp things
self.lfp_plant.turn_off()
self.targets[0].hide()
def _start_lfp_target(self):
#only 1 target:
target = self.targets[0] #LFP target
self.target_location_lfp = self.lfp_targ #LFP target
target.move_to_position(self.target_location_lfp)
target.cue_trial_start()
self.lfp_plant.turn_on()
def _start_lfp_hold(self):
#make next target visible unless this is the final target in the trial
pass
def _start_hold_penalty(self):
#hide targets
for target in self.targets:
target.hide()
self.tries += 1
self.target_index = -1
#Turn off lfp things
self.lfp_plant.turn_off()
self.targets[0].hide()
def _end_origin(self):
self.targets[1].cue_trial_end_success()
def _test_enter_origin(self, ts):
cursor_pos = self.mc_plant.get_endpoint_pos()
d = np.linalg.norm(cursor_pos - self.target_location_mc)
return d <= (self.mc_target_radius - self.mc_cursor_radius)
# def _test_origin_timeout(self, ts):
# return ts>self.timeout_time
def _test_leave_origin(self, ts):
if self.manual_control_type == 'joystick':
if hasattr(self,'touch'):
if self.touch <0.5:
self.last_touch_zero_event = time.time()
return True
cursor_pos = self.mc_plant.get_endpoint_pos()
d = np.linalg.norm(cursor_pos - self.target_location_mc)
return d > (self.mc_target_radius - self.mc_cursor_radius)
def _test_origin_hold_complete(self,ts):
return ts>=self.origin_hold_time
# def _test_enter_lfp_target(self, ts):
# '''
# return true if the distance between center of cursor and target is smaller than the cursor radius
# '''
# cursor_pos = self.lfp_plant.get_endpoint_pos()
# cursor_pos = [cursor_pos[0], cursor_pos[2]]
# targ_loc = np.array([self.target_location_lfp[0], self.target_location_lfp[2]])
# d = np.linalg.norm(cursor_pos - targ_loc)
# return d <= (self.lfp_target_rad - self.lfp_cursor_rad)
# def _test_leave_early(self, ts):
# '''
# return true if cursor moves outside the exit radius
# '''
# cursor_pos = self.lfp_plant.get_endpoint_pos()
# d = np.linalg.norm(cursor_pos - self.target_location_lfp)
# rad = self.lfp_target_rad - self.lfp_cursor_rad
# return d > rad
def _test_hold_penalty_end(self, ts):
return ts>self.hold_penalty_time
def _end_lfp_hold(self):
# change current target color to green
self.targets[0].cue_trial_end_success()
def move_plant(self):
if self.state in self.lfp_cursor_on:
feature_data = self.get_features()
# Save the "neural features" (e.g. spike counts vector) to HDF file
for key, val in feature_data.items():
self.task_data[key] = val
Bu = None
assist_weight = 0
target_state = np.zeros([self.decoder.n_states, self.decoder.n_subbins])
## Run the decoder
neural_features = feature_data[self.extractor.feature_type]
self.call_decoder(neural_features, target_state, Bu=Bu, assist_level=assist_weight, feature_type=self.extractor.feature_type)
## Drive the plant to the decoded state, if permitted by the constraints of the plant
self.lfp_plant.drive(self.decoder)
self.task_data['decoder_state'] = decoder_state = self.decoder.get_state(shape=(-1,1))
#return decoder_state
#Sets the plant configuration based on motiontracker data. For manual control, uses
#motiontracker data. If no motiontracker data available, returns None'''
#get data from motion tracker- take average of all data points since last poll
if self.state in self.no_plant_update_states:
pt = np.array([0, 0, 0])
print 'no update'
else:
if self.manual_control_type == 'motiondata':
pt = self.motiondata.get()
if len(pt) > 0:
pt = pt[:, self.marker_num, :]
conds = pt[:, 3]
inds = np.nonzero((conds>=0) & (conds!=4))[0]
if len(inds) > 0:
pt = pt[inds,:3]
#scale actual movement to desired amount of screen movement
pt = pt.mean(0) * self.scale_factor
#Set y coordinate to 0 for 2D tasks
if self.limit2d:
#pt[1] = 0
pt[2] = pt[1].copy()
pt[1] = 0
pt[1] = pt[1]*2
# Return cursor location
self.no_data_count = 0
pt = pt * mm_per_cm #self.convert_to_cm(pt)
else: #if no usable data
self.no_data_count += 1
pt = None
else: #if no new data
self.no_data_count +=1
pt = None
elif self.manual_control_type == 'joystick':
pt = self.joystick.get()
#if touch sensor on:
try:
self.touch = pt[-1][0][2]
except:
pass
if len(pt) > 0:
pt = pt[-1][0]
pt[0]=1-pt[0]; #Switch L / R axes
calib = [0.497,0.517] #Sometimes zero point is subject to drift this is the value of the incoming joystick when at 'rest'
if self.joystick_method==0:
#pt = pt[-1][0]
#pt[0]=1-pt[0]; #Switch L / R axes
#calib = [0.497,0.517] #Sometimes zero point is subject to drift this is the value of the incoming joystick when at 'rest'
# calib = [ 0.487, 0. ]
pos = np.array([(pt[0]-calib[0]), 0, calib[1]-pt[1]])
pos[0] = pos[0]*36
pos[2] = pos[2]*24
self.current_pt = pos
elif self.joystick_method==1:
vel=np.array([(pt[0]-calib[0]), 0, calib[1]-pt[1]])
epsilon = 2*(10**-2) #Define epsilon to stabilize cursor movement
if sum((vel)**2) > epsilon:
self.current_pt=self.last_pt+self.joystick_speed*vel*(1/60) #60 Hz update rate, dt = 1/60
else:
self.current_pt = self.last_pt
if self.current_pt[0] < -25: self.current_pt[0] = -25
if self.current_pt[0] > 25: self.current_pt[0] = 25
if self.current_pt[-1] < -14: self.current_pt[-1] = -14
if self.current_pt[-1] > 14: self.current_pt[-1] = 14
pt = self.current_pt
#self.plant.set_endpoint_pos(self.current_pt)
self.last_pt = self.current_pt.copy()
elif self.manual_control_type == None:
pt = None
try:
pt0 = self.motiondata.get()
self.manual_control_type='motiondata'
except:
print 'not motiondata'
try:
pt0 = self.joystick.get()
self.manual_control_type = 'joystick'
except:
print 'not joystick data'
# Set the plant's endpoint to the position determined by the motiontracker, unless there is no data available
if self.manual_control_type is not None:
if pt is not None and len(pt)>0:
self.mc_plant.set_endpoint_pos(pt)
@staticmethod
def lfp_mod_4targ_plus_mc_orig(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8):
'''
See lfp_mod_4targ for lfp target explanation
'''
new_zero = (boundaries[3]+boundaries[2]) / 2.
new_scale = (boundaries[3] - boundaries[2]) / 2.
kin_targs = np.array([-0.40625, 0. , 0.40625, 0.8125 ])
lfp_targ_y = (new_scale*kin_targs) + new_zero
for i in range(nblocks):
temp = lfp_targ_y.copy()
np.random.shuffle(temp)
if i==0:
z = temp.copy()
else:
z = np.hstack((z, temp))
#Fixed X axis:
x = np.tile(xaxis,(nblocks*4))
y = np.zeros(nblocks*4)
lfp = np.vstack([x, y, z]).T
origin = np.zeros(( lfp.shape ))
it = iter([dict(lfp=lfp[i,:], origin=origin[i,:]) for i in range(lfp.shape[0])])
return it
class ApproachAvoidanceTask(Sequence, Window):
'''
This is for a free-choice task with two targets (left and right) presented to choose from.
The position of the targets may change along the x-axis, according to the target generator,
and each target has a varying probability of reward, also according to the target generator.
The code as it is written is for a joystick.
Notes: want target_index to only write once per trial. if so, can make instructed trials random. else, make new state for instructed trial.
'''
background = (0,0,0,1)
shoulder_anchor = np.array([2., 0., -15.]) # Coordinates of shoulder anchor point on screen
arm_visible = traits.Bool(True, desc='Specifies whether entire arm is displayed or just endpoint')
cursor_radius = traits.Float(.5, desc="Radius of cursor")
cursor_color = (.5,0,.5,1)
joystick_method = traits.Float(1,desc="1: Normal velocity, 0: Position control")
joystick_speed = traits.Float(20, desc="Speed of cursor")
plant_type_options = plantlist.keys()
plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys())
starting_pos = (5, 0, 5)
# window_size = (1280*2, 1024)
window_size = traits.Tuple((1366*2, 768), desc='window size')
status = dict(
#wait = dict(start_trial="target", stop=None),
wait = dict(start_trial="center", stop=None),
center = dict(enter_center="hold_center", timeout="timeout_penalty", stop=None),
hold_center = dict(leave_early_center = "hold_penalty",hold_center_complete="target", timeout="timeout_penalty", stop=None),
target = dict(enter_targetL="hold_targetL", enter_targetH = "hold_targetH", timeout="timeout_penalty", stop=None),
hold_targetR = dict(leave_early_R="hold_penalty", hold_complete="targ_transition"),
hold_targetL = dict(leave_early_L="hold_penalty", hold_complete="targ_transition"),
targ_transition = dict(trial_complete="check_reward",trial_abort="wait", trial_incomplete="center"),
check_reward = dict(avoid="reward",approach="reward_and_airpuff"),
timeout_penalty = dict(timeout_penalty_end="targ_transition"),
hold_penalty = dict(hold_penalty_end="targ_transition"),
reward = dict(reward_end="wait"),
reward_and_airpuff = dict(reward_and_airpuff_end="wait"),
)
#
target_color = (.5,1,.5,0)
#initial state
state = "wait"
#create settable traits
reward_time_avoid = traits.Float(.2, desc="Length of juice reward for avoid decision")
reward_time_approach_min = traits.Float(.2, desc="Min length of juice for approach decision")
reward_time_approach_max = traits.Float(.8, desc="Max length of juice for approach decision")
target_radius = traits.Float(1.5, desc="Radius of targets in cm")
block_length = traits.Float(100, desc="Number of trials per block")
hold_time = traits.Float(.5, desc="Length of hold required at targets")
hold_penalty_time = traits.Float(1, desc="Length of penalty time for target hold error")
timeout_time = traits.Float(10, desc="Time allowed to go between targets")
timeout_penalty_time = traits.Float(1, desc="Length of penalty time for timeout error")
max_attempts = traits.Int(10, desc='The number of attempts at a target before\
skipping to the next one')
session_length = traits.Float(0, desc="Time until task automatically stops. Length of 0 means no auto stop.")
marker_num = traits.Int(14, desc="The index of the motiontracker marker to use for cursor position")
arm_hide_rate = traits.Float(0.0, desc='If the arm is visible, specifies a percentage of trials where it will be hidden')
target_index = 0 # Helper variable to keep track of whether trial is instructed (1 = 1 choice) or free-choice (2 = 2 choices)
target_selected = 'L' # Helper variable to indicate which target was selected
tries = 0 # Helper variable to keep track of the number of failed attempts at a given trial.
timedout = False # Helper variable to keep track if transitioning from timeout_penalty
reward_counter = 0.0
cursor_visible = False # Determines when to hide the cursor.
no_data_count = 0 # Counter for number of missing data frames in a row
scale_factor = 3.0 #scale factor for converting hand movement to screen movement (1cm hand movement = 3.5cm cursor movement)
starting_dist = 10.0 # starting distance from center target
#color_targets = np.random.randint(2)
color_targets = 1 # 0: yellow low, blue high; 1: blue low, yellow high
stopped_center_hold = False #keep track if center hold was released early
limit2d = 1
color1 = target_colors['purple'] # approach color
color2 = target_colors['lightsteelblue'] # avoid color
reward_color = target_colors['green'] # color of reward bar
airpuff_color = target_colors['red'] # color of airpuff bar
sequence_generators = ['colored_targets_with_probabilistic_reward','block_probabilistic_reward','colored_targets_with_randomwalk_reward','randomwalk_probabilistic_reward']
def __init__(self, *args, **kwargs):
super(ApproachAvoidanceTask, self).__init__(*args, **kwargs)
self.cursor_visible = True
# Add graphics models for the plant and targets to the window
self.plant = plantlist[self.plant_type]
self.plant_vis_prev = True
# Add graphics models for the plant and targets to the window
if hasattr(self.plant, 'graphics_models'):
for model in self.plant.graphics_models:
self.add_model(model)
self.current_pt=np.zeros([3]) #keep track of current pt
self.last_pt=np.zeros([3]) #kee
## Declare cursor
#self.dtype.append(('cursor', 'f8', (3,)))
if 0: #hasattr(self.arm, 'endpt_cursor'):
self.cursor = self.arm.endpt_cursor
else:
self.cursor = Sphere(radius=self.cursor_radius, color=self.cursor_color)
self.add_model(self.cursor)
self.cursor.translate(*self.get_arm_endpoint(), reset=True)
## Instantiate the targets. Target 1 is center target, Target H is target with high probability of reward, Target L is target with low probability of reward.
self.target1 = Sphere(radius=self.target_radius, color=self.target_color) # center target
self.add_model(self.target1)
self.targetR = Sphere(radius=self.target_radius, color=self.target_color) # left target
self.add_model(self.targetH)
self.targetL = Sphere(radius=self.target_radius, color=self.target_color) # right target
self.add_model(self.targetL)
###STOPPED HERE: should define Rect target here and then adapt length during task. Also,
### be sure to change all targetH instantiations to targetR.
# Initialize target location variable.
self.target_location1 = np.array([0,0,0])
self.target_locationR = np.array([-self.starting_dist,0,0])
self.target_locationL = np.array([self.starting_dist,0,0])
self.target1.translate(*self.target_location1, reset=True)
self.targetH.translate(*self.target_locationR, reset=True)
self.targetL.translate(*self.target_locationL, reset=True)
# Initialize colors for high probability and low probability target. Color will not change.
self.targetH.color = self.color_targets*self.color1 + (1 - self.color_targets)*self.color2 # high is magenta if color_targets = 1, juicyorange otherwise
self.targetL.color = (1 - self.color_targets)*self.color1 + self.color_targets*self.color2
#set target colors
self.target1.color = (1,0,0,.5) # center target red
# Initialize target location variable
self.target_location = np.array([0, 0, 0])
# Declare any plant attributes which must be saved to the HDF file at the _cycle rate
for attr in self.plant.hdf_attrs:
self.add_dtype(*attr)
def init(self):
self.add_dtype('targetR', 'f8', (3,))
self.add_dtype('targetL','f8', (3,))
self.add_dtype('reward_scheduleR','f8', (1,))
self.add_dtype('reward_scheduleL','f8', (1,))
self.add_dtype('target_index', 'i', (1,))
super(ApproachAvoidanceTask, self).init()
self.trial_allocation = np.zeros(1000)
def _cycle(self):
''' Calls any update functions necessary and redraws screen. Runs 60x per second. '''
## Run graphics commands to show/hide the arm if the visibility has changed
if self.plant_type != 'cursor_14x14':
if self.arm_visible != self.arm_vis_prev:
self.arm_vis_prev = self.arm_visible
self.show_object(self.arm, show=self.arm_visible)
self.move_arm()
#self.move_plant()
## Save plant status to HDF file
plant_data = self.plant.get_data_to_save()
for key in plant_data:
self.task_data[key] = plant_data[key]
self.update_cursor()
if self.plant_type != 'cursor_14x14':
self.task_data['joint_angles'] = self.get_arm_joints()
super(ApproachAvoidanceTask, self)._cycle()
## Plant functions
def get_cursor_location(self):
# arm returns it's position as if it was anchored at the origin, so have to translate it to the correct place
return self.get_arm_endpoint()
def get_arm_endpoint(self):
return self.plant.get_endpoint_pos()
def set_arm_endpoint(self, pt, **kwargs):
self.plant.set_endpoint_pos(pt, **kwargs)
def set_arm_joints(self, angles):
self.arm.set_intrinsic_coordinates(angles)
def get_arm_joints(self):
return self.arm.get_intrinsic_coordinates()
def update_cursor(self):
'''
Update the cursor's location and visibility status.
'''
pt = self.get_cursor_location()
self.update_cursor_visibility()
if pt is not None:
self.move_cursor(pt)
def move_cursor(self, pt):
''' Move the cursor object to the specified 3D location. '''
# if not hasattr(self.arm, 'endpt_cursor'):
self.cursor.translate(*pt[:3],reset=True)
##
##### HELPER AND UPDATE FUNCTIONS ####
def move_arm(self):
''' Returns the 3D coordinates of the cursor. For manual control, uses
joystick data. If no joystick data available, returns None'''
pt = self.joystick.get()
if len(pt) > 0:
pt = pt[-1][0]
pt[0]=1-pt[0]; #Switch L / R axes
calib = [0.497,0.517] #Sometimes zero point is subject to drift this is the value of the incoming joystick when at 'rest'
if self.joystick_method==0:
pos = np.array([(pt[0]-calib[0]), 0, calib[1]-pt[1]])
pos[0] = pos[0]*36
pos[2] = pos[2]*24
self.current_pt = pos
elif self.joystick_method==1:
vel=np.array([(pt[0]-calib[0]), 0, calib[1]-pt[1]])
epsilon = 2*(10**-2) #Define epsilon to stabilize cursor movement
if sum((vel)**2) > epsilon:
self.current_pt=self.last_pt+self.joystick_speed*vel*(1/60) #60 Hz update rate, dt = 1/60
else:
self.current_pt = self.last_pt
if self.current_pt[0] < -25: self.current_pt[0] = -25
if self.current_pt[0] > 25: self.current_pt[0] = 25
if self.current_pt[-1] < -14: self.current_pt[-1] = -14
if self.current_pt[-1] > 14: self.current_pt[-1] = 14
self.set_arm_endpoint(self.current_pt)
self.last_pt = self.current_pt.copy()
def convert_to_cm(self, val):
return val/10.0
def update_cursor_visibility(self):
''' Update cursor visible flag to hide cursor if there has been no good data for more than 3 frames in a row'''
prev = self.cursor_visible
if self.no_data_count < 3:
self.cursor_visible = True
if prev != self.cursor_visible:
self.show_object(self.cursor, show=True)
self.requeue()
else:
self.cursor_visible = False
if prev != self.cursor_visible:
self.show_object(self.cursor, show=False)
self.requeue()
def calc_n_successfultrials(self):
trialendtimes = np.array([state[1] for state in self.state_log if state[0]=='check_reward'])
return len(trialendtimes)
def calc_n_rewards(self):
rewardtimes = np.array([state[1] for state in self.state_log if state[0]=='reward'])
return len(rewardtimes)
def calc_trial_num(self):
'''Calculates the current trial count: completed + aborted trials'''
trialtimes = [state[1] for state in self.state_log if state[0] in ['wait']]
return len(trialtimes)-1
def calc_targetH_num(self):
'''Calculates the total number of times the high-value target was selected'''
trialtimes = [state[1] for state in self.state_log if state[0] in ['hold_targetH']]
return len(trialtimes) - 1
def calc_rewards_per_min(self, window):
'''Calculates the Rewards/min for the most recent window of specified number of seconds in the past'''
rewardtimes = np.array([state[1] for state in self.state_log if state[0]=='reward'])
if (self.get_time() - self.task_start_time) < window:
divideby = (self.get_time() - self.task_start_time)/sec_per_min
else:
divideby = window/sec_per_min
return np.sum(rewardtimes >= (self.get_time() - window))/divideby
def calc_success_rate(self, window):
'''Calculates the rewarded trials/initiated trials for the most recent window of specified length in sec'''
trialtimes = np.array([state[1] for state in self.state_log if state[0] in ['reward', 'timeout_penalty', 'hold_penalty']])
rewardtimes = np.array([state[1] for state in self.state_log if state[0]=='reward'])
if len(trialtimes) == 0:
return 0.0
else:
return float(np.sum(rewardtimes >= (self.get_time() - window)))/np.sum(trialtimes >= (self.get_time() - window))
def update_report_stats(self):
'''Function to update any relevant report stats for the task. Values are saved in self.reportstats,
an ordered dictionary. Keys are strings that will be displayed as the label for the stat in the web interface,
values can be numbers or strings. Called every time task state changes.'''
super(ApproachAvoidanceTask, self).update_report_stats()
self.reportstats['Trial #'] = self.calc_trial_num()
self.reportstats['Reward/min'] = np.round(self.calc_rewards_per_min(120),decimals=2)
self.reportstats['High-value target selections'] = self.calc_targetH_num()
#self.reportstats['Success rate'] = str(np.round(self.calc_success_rate(120)*100.0,decimals=2)) + '%'
start_time = self.state_log[0][1]
rewardtimes=np.array([state[1] for state in self.state_log if state[0]=='reward'])
if len(rewardtimes):
rt = rewardtimes[-1]-start_time
else:
rt= np.float64("0.0")
sec = str(np.int(np.mod(rt,60)))
if len(sec) < 2:
sec = '0'+sec
self.reportstats['Time Of Last Reward'] = str(np.int(np.floor(rt/60))) + ':' + sec
#### TEST FUNCTIONS ####
def _test_enter_center(self, ts):
#return true if the distance between center of cursor and target is smaller than the cursor radius
d = np.sqrt((self.cursor.xfm.move[0]-self.target_location1[0])**2 + (self.cursor.xfm.move[1]-self.target_location1[1])**2 + (self.cursor.xfm.move[2]-self.target_location1[2])**2)
#print 'TARGET SELECTED', self.target_selected
return d <= self.target_radius - self.cursor_radius
def _test_enter_targetL(self, ts):
if self.target_index == 1 and self.LH_target_on[0]==0:
#return false if instructed trial and this target is not on
return False
else:
#return true if the distance between center of cursor and target is smaller than the cursor radius
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationL[0])**2 + (self.cursor.xfm.move[1]-self.target_locationL[1])**2 + (self.cursor.xfm.move[2]-self.target_locationL[2])**2)
self.target_selected = 'L'
#print 'TARGET SELECTED', self.target_selected
return d <= self.target_radius - self.cursor_radius
def _test_enter_targetH(self, ts):
if self.target_index ==1 and self.LH_target_on[1]==0:
return False
else:
#return true if the distance between center of cursor and target is smaller than the cursor radius
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationH[0])**2 + (self.cursor.xfm.move[1]-self.target_locationH[1])**2 + (self.cursor.xfm.move[2]-self.target_locationH[2])**2)
self.target_selected = 'H'
#print 'TARGET SELECTED', self.target_selected
return d <= self.target_radius - self.cursor_radius
def _test_leave_early_center(self, ts):
# return true if cursor moves outside the exit radius (gives a bit of slack around the edge of target once cursor is inside)
d = np.sqrt((self.cursor.xfm.move[0]-self.target_location1[0])**2 + (self.cursor.xfm.move[1]-self.target_location1[1])**2 + (self.cursor.xfm.move[2]-self.target_location1[2])**2)
rad = self.target_radius - self.cursor_radius
return d > rad
def _test_leave_early_L(self, ts):
# return true if cursor moves outside the exit radius (gives a bit of slack around the edge of target once cursor is inside)
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationL[0])**2 + (self.cursor.xfm.move[1]-self.target_locationL[1])**2 + (self.cursor.xfm.move[2]-self.target_locationL[2])**2)
rad = self.target_radius - self.cursor_radius
return d > rad
def _test_leave_early_H(self, ts):
# return true if cursor moves outside the exit radius (gives a bit of slack around the edge of target once cursor is inside)
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationH[0])**2 + (self.cursor.xfm.move[1]-self.target_locationH[1])**2 + (self.cursor.xfm.move[2]-self.target_locationH[2])**2)
rad = self.target_radius - self.cursor_radius
return d > rad
def _test_hold_center_complete(self, ts):
return ts>=self.hold_time
def _test_hold_complete(self, ts):
return ts>=self.hold_time
def _test_timeout(self, ts):
return ts>self.timeout_time
def _test_timeout_penalty_end(self, ts):
return ts>self.timeout_penalty_time
def _test_hold_penalty_end(self, ts):
return ts>self.hold_penalty_time
def _test_trial_complete(self, ts):
#return self.target_index==self.chain_length-1
return not self.timedout
def _test_trial_incomplete(self, ts):
return (not self._test_trial_complete(ts)) and (self.tries<self.max_attempts)
def _test_trial_abort(self, ts):
return (not self._test_trial_complete(ts)) and (self.tries==self.max_attempts)
def _test_yes_reward(self,ts):
if self.target_selected == 'H':
#reward_assigned = self.targs[0,1]
reward_assigned = self.rewardH
else:
#reward_assigned = self.targs[1,1]
reward_assigned = self.rewardL
if self.reward_SmallLarge==1:
self.reward_time = reward_assigned*self.reward_time_large + (1 - reward_assigned)*self.reward_time_small # update reward time if using Small/large schedule
reward_assigned = 1 # always rewarded
return bool(reward_assigned)
def _test_no_reward(self,ts):
if self.target_selected == 'H':
#reward_assigned = self.targs[0,1]
reward_assigned = self.rewardH
else:
#reward_assigned = self.targs[1,1]
reward_assigned = self.rewardL
if self.reward_SmallLarge==True:
self.reward_time = reward_assigned*self.reward_time_large + (1 - reward_assigned)*self.reward_time_small # update reward time if using Small/large schedule
reward_assigned = 1 # always rewarded
return bool(not reward_assigned)
def _test_reward_end(self, ts):
time_ended = (ts > self.reward_time)
self.reward_counter = self.reward_counter + 1
return time_ended
def _test_stop(self, ts):
if self.session_length > 0 and (time.time() - self.task_start_time) > self.session_length:
self.end_task()
return self.stop
#### STATE FUNCTIONS ####
def show_object(self, obj, show=False):
'''
Show or hide an object
'''
if show:
obj.attach()
else:
obj.detach()
self.requeue()
def _start_wait(self):
super(ApproachAvoidanceTask, self)._start_wait()
self.tries = 0
self.target_index = 0 # indicator for instructed or free-choice trial
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
#get target positions and reward assignments for this trial
self.targs = self.next_trial
if self.plant_type != 'cursor_14x14' and np.random.rand() < self.arm_hide_rate:
self.arm_visible = False
else:
self.arm_visible = True
#self.chain_length = self.targs.shape[0] #Number of sequential targets in a single trial
#self.task_data['target'] = self.target_locationH.copy()
assign_reward = np.random.randint(0,100,size=2)
self.rewardH = np.greater(self.targs[0,1],assign_reward[0])
#print 'high value target reward prob', self.targs[0,1]
self.rewardL = np.greater(self.targs[1,1],assign_reward[1])
#print 'TARGET GENERATOR', self.targs[0,]
self.task_data['targetH'] = self.targs[0,].copy()
self.task_data['reward_scheduleH'] = self.rewardH.copy()
self.task_data['targetL'] = self.targs[1,].copy()
self.task_data['reward_scheduleL'] = self.rewardL.copy()
self.requeue()
def _start_center(self):
#self.target_index += 1
self.show_object(self.target1, True)
self.show_object(self.cursor, True)
# Third argument in self.targs determines if target is on left or right
# First argument in self.targs determines if location is offset to farther distances
offsetH = (2*self.targs[0,2] - 1)*(self.starting_dist + self.location_offset_allowed*self.targs[0,0]*4.0)
moveH = np.array([offsetH,0,0])
offsetL = (2*self.targs[1,2] - 1)*(self.starting_dist + self.location_offset_allowed*self.targs[1,0]*4.0)
moveL = np.array([offsetL,0,0])
self.targetL.translate(*moveL, reset=True)
#self.targetL.move_to_position(*moveL, reset=True)
##self.targetL.translate(*self.targs[self.target_index], reset=True)
self.show_object(self.targetL, True)
self.target_locationL = self.targetL.xfm.move
self.targetH.translate(*moveH, reset=True)
#self.targetR.move_to_position(*moveR, reset=True)
##self.targetR.translate(*self.targs[self.target_index], reset=True)
self.show_object(self.targetH, True)
self.target_locationH = self.targetH.xfm.move
# Insert instructed trials within free choice trials
if self.trial_allocation[self.calc_trial_num()] == 1:
#if (self.calc_trial_num() % 10) < (self.percentage_instructed_trials/10):
self.target_index = 1 # instructed trial
leftright_coinflip = np.random.randint(0,2)
if leftright_coinflip == 0:
self.show_object(self.targetL, False)
self.LH_target_on = (0, 1)
else:
self.show_object(self.targetH, False)
self.LR_coinflip = 0
self.LH_target_on = (1, 0)
else:
self.target_index = 2 # free-choice trial
self.cursor_visible = True
self.task_data['target_index'] = self.target_index
self.requeue()
def _start_target(self):
#self.target_index += 1
#move targets to current location and set location attribute. Target1 (center target) position does not change.
self.show_object(self.target1, False)
#self.target_location1 = self.target1.xfm.move
self.show_object(self.cursor, True)
self.update_cursor()
self.requeue()
def _start_hold_center(self):
self.show_object(self.target1, True)
self.timedout = False
self.requeue()
def _start_hold_targetL(self):
#make next target visible unless this is the final target in the trial
#if 1 < self.chain_length:
#self.targetL.translate(*self.targs[self.target_index+1], reset=True)
# self.show_object(self.targetL, True)
# self.requeue()
self.show_object(self.targetL, True)
self.timedout = False
self.requeue()
def _start_hold_targetH(self):
#make next target visible unless this is the final target in the trial
#if 1 < self.chain_length:
#self.targetR.translate(*self.targs[self.target_index+1], reset=True)
# self.show_object(self.targetR, True)
# self.requeue()
self.show_object(self.targetH, True)
self.timedout = False
self.requeue()
def _end_hold_center(self):
self.target1.radius = 0.7*self.target_radius # color target green
def _end_hold_targetL(self):
self.targetL.color = (0,1,0,0.5) # color target green
def _end_hold_targetH(self):
self.targetH.color = (0,1,0,0.5) # color target green
def _start_hold_penalty(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.timedout = True
self.requeue()
self.tries += 1
#self.target_index = -1
def _start_timeout_penalty(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.timedout = True
self.requeue()
self.tries += 1
#self.target_index = -1
def _start_targ_transition(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.requeue()
def _start_check_reward(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.requeue()
def _start_reward(self):
#super(ApproachAvoidanceTask, self)._start_reward()
if self.target_selected == 'L':
self.show_object(self.targetL, True)
#reward_assigned = self.targs[1,1]
else:
self.show_object(self.targetH, True)
#reward_assigned = self.targs[0,1]
#self.reward_counter = self.reward_counter + float(reward_assigned)
self.requeue()
@staticmethod
def colored_targets_with_probabilistic_reward(length=1000, boundaries=(-18,18,-10,10,-15,15),reward_high_prob=80,reward_low_prob=40):
"""
Generator should return array of ntrials x 2 x 3. The second dimension is for each target.
For example, first is the target with high probability of reward, and the second
entry is for the target with low probability of reward. The third dimension holds three variables indicating
position offset (yes/no), reward probability (fixed in this case), and location (binary returned where the
ouput indicates either left or right).
UPDATE: CHANGED SO THAT THE SECOND DIMENSION CARRIES THE REWARD PROBABILITY RATHER THAN THE REWARD SCHEDULE
"""
position_offsetH = np.random.randint(2,size=(1,length))
position_offsetL = np.random.randint(2,size=(1,length))
location_int = np.random.randint(2,size=(1,length))
# coin flips for reward schedules, want this to be elementwise comparison
#assign_rewardH = np.random.randint(0,100,size=(1,length))
#assign_rewardL = np.random.randint(0,100,size=(1,length))
high_prob = reward_high_prob*np.ones((1,length))
low_prob = reward_low_prob*np.ones((1,length))
#reward_high = np.greater(high_prob,assign_rewardH)
#reward_low = np.greater(low_prob,assign_rewardL)
pairs = np.zeros([length,2,3])
pairs[:,0,0] = position_offsetH
#pairs[:,0,1] = reward_high
pairs[:,0,1] = high_prob
pairs[:,0,2] = location_int
pairs[:,1,0] = position_offsetL
#pairs[:,1,1] = reward_low
pairs[:,1,1] = low_prob
pairs[:,1,2] = 1 - location_int
return pairs
@staticmethod
def block_probabilistic_reward(length=1000, boundaries=(-18,18,-10,10,-15,15),reward_high_prob=80,reward_low_prob=40):
pairs = colored_targets_with_probabilistic_reward(length=length, boundaries=boundaries,reward_high_prob=reward_high_prob,reward_low_prob=reward_low_prob)
return pairs
@staticmethod
def colored_targets_with_randomwalk_reward(length=1000,reward_high_prob=80,reward_low_prob=40,reward_high_span = 20, reward_low_span = 20,step_size_mean = 0, step_size_var = 1):
"""
Generator should return array of ntrials x 2 x 3. The second dimension is for each target.
For example, first is the target with high probability of reward, and the second
entry is for the target with low probability of reward. The third dimension holds three variables indicating
position offset (yes/no), reward probability, and location (binary returned where the
ouput indicates either left or right). The variables reward_high_span and reward_low_span indicate the width
of the range that the high or low reward probability are allowed to span respectively, e.g. if reward_high_prob
is 80 and reward_high_span is 20, then the reward probability for the high value target will be bounded
between 60 and 100 percent.
"""
position_offsetH = np.random.randint(2,size=(1,length))
position_offsetL = np.random.randint(2,size=(1,length))
location_int = np.random.randint(2,size=(1,length))
# define variables for increments: amount of increment and in which direction (i.e. increasing or decreasing)
assign_rewardH = np.random.randn(1,length)
assign_rewardL = np.random.randn(1,length)
assign_rewardH_direction = np.random.randn(1,length)
assign_rewardL_direction = np.random.randn(1,length)
r_0_high = reward_high_prob
r_0_low = reward_low_prob
r_lowerbound_high = r_0_high - (reward_high_span/2)
r_upperbound_high = r_0_high + (reward_high_span/2)
r_lowerbound_low = r_0_low - (reward_low_span/2)
r_upperbound_low = r_0_low + (reward_low_span/2)
reward_high = np.zeros(length)
reward_low = np.zeros(length)
reward_high[0] = r_0_high
reward_low[0] = r_0_low
eps_high = assign_rewardH*step_size_mean + [2*(assign_rewardH_direction > 0) - 1]*step_size_var
eps_low = assign_rewardL*step_size_mean + [2*(assign_rewardL_direction > 0) - 1]*step_size_var
eps_high = eps_high.ravel()
eps_low = eps_low.ravel()
for i in range(1,length):
'''
assign_rewardH_direction = np.random.randn(1)
assign_rewardL_direction = np.random.randn(1)
assign_rewardH = np.random.randn(1)
if assign_rewardH_direction[i-1,] < 0:
eps_high = step_size_mean*assign_rewardH[i-1] - step_size_var
else:
eps_high = step_size_mean*assign_rewardH[i-1] + step_size_var
if assign_rewardL_direction[i] < 0:
eps_low = step_size_mean*assign_rewardL[i] - step_size_var
else:
eps_low = step_size_mean*assign_rewardL[i] + step_size_var
'''
reward_high[i] = reward_high[i-1] + eps_high[i-1]
reward_low[i] = reward_low[i-1] + eps_low[i-1]
reward_high[i] = (r_lowerbound_high < reward_high[i] < r_upperbound_high)*reward_high[i] + (r_lowerbound_high > reward_high[i])*(r_lowerbound_high+ eps_high[i-1]) + (r_upperbound_high < reward_high[i])*(r_upperbound_high - eps_high[i-1])
reward_low[i] = (r_lowerbound_low < reward_low[i] < r_upperbound_low)*reward_low[i] + (r_lowerbound_low > reward_low[i])*(r_lowerbound_low+ eps_low[i-1]) + (r_upperbound_low < reward_low[i])*(r_upperbound_low - eps_low[i-1])
pairs = np.zeros([length,2,3])
pairs[:,0,0] = position_offsetH
pairs[:,0,1] = reward_high
pairs[:,0,2] = location_int
pairs[:,1,0] = position_offsetL
pairs[:,1,1] = reward_low
pairs[:,1,2] = 1 - location_int
return pairs
@staticmethod
def randomwalk_probabilistic_reward(length=1000,reward_high_prob=80,reward_low_prob=40,reward_high_span = 20, reward_low_span = 20,step_size_mean = 0, step_size_var = 1):
pairs = colored_targets_with_randomwalk_reward(length=length,reward_high_prob=reward_high_prob,reward_low_prob=reward_low_prob,reward_high_span = reward_high_span, reward_low_span = reward_low_span,step_size_mean = step_size_mean, step_size_var = step_size_var)
return pairs