def _grab_RNN(self, initial_states): '''Creates objects for interfacing with the RNN. These objects include 1) the optimization variables (initialized to the user-specified initial_states) which will, after optimization, contain fixed points of the RNN, and 2) hooks into those optimization variables that are required for building the TF graph. Args: initial_states: Either an [n_inits x n_dims] numpy array or an LSTMStateTuple with initial_states.c and initial_states.h as [n_inits x n_dims/2] numpy arrays. These data specify the initial states of the RNN, from which the optimization will search for fixed points. The choice of type must be consistent with state type of rnn_cell. Returns: x: An [n_inits x n_dims] tf.Variable (the optimization variable) representing RNN states, initialized to the values in initial_states. If the RNN is an LSTM, n_dims represents the concatenated hidden and cell states. F: An [n_inits x n_dims] tf op representing the state transition function of the RNN applied to x. states: Contains the same data as in x, but formatted to interface with self.rnn_cell (e.g., formatted as LSTMStateTuple if rnn_cell is a LSTMCell) new_states: Contains the same data as in F, but formatted to interface with self.rnn_cell ''' if self.is_lstm: # [1 x (2*n_dims)] c_h_init = tf_utils.convert_from_LSTMStateTuple(initial_states) # [1 x (2*n_dims)] x = tf.Variable(c_h_init, dtype=tf.float32) states = tf_utils.convert_to_LSTMStateTuple(x) else: x = tf.Variable(initial_states, dtype=tf.float32) states = x n_inits = x.shape[0] tiled_inputs = np.tile(self.inputs, [n_inits, 1]) inputs_tf = tf.constant(tiled_inputs, dtype=tf.float32) output, new_states = self.rnn_cell(inputs_tf, states) if self.is_lstm: # [1 x (2*n_dims)] F = tf_utils.convert_from_LSTMStateTuple(new_states) else: F = new_states init = tf.variables_initializer(var_list=[x]) self.session.run(init) return x, F, states, new_states
def _compute_multiple_jacobians_np(self, fps): '''Computes the Jacobian of the RNN state transition function. Args: fps: A FixedPoints object containing the RNN states (fps.xstar) and inputs (fps.inputs) at which to compute the Jacobians. Returns: J_np: An [n x n_states x n_states] numpy array containing the Jacobian of the RNN state transition function at the states specified in fps, given the inputs in fps. ''' inputs_np = fps.inputs if self.is_lstm: states_np = tf_utils.convert_to_LSTMStateTuple(fps.xstar) else: states_np = fps.xstar x_tf, F_tf = self._grab_RNN(states_np, inputs_np) try: J_tf = pfor.batch_jacobian(F_tf, x_tf) except absl.flags._exceptions.UnparsedFlagAccessError: J_tf = pfor.batch_jacobian(F_tf, x_tf, use_pfor=False) J_np = self.session.run(J_tf) return J_np
def sample_states(self, state_traj, n_inits, noise_scale=0.0, rng=npr.RandomState(0)): '''Draws random samples from trajectories of the RNN state. Samples can optionally be corrupted by independent and identically distributed (IID) Gaussian noise. These samples are intended to be used as initial states for fixed point optimizations. Args: state_traj: [n_batch x n_time x n_states] numpy array or LSTMStateTuple with .c and .h as [n_batch x n_time x n_states] numpy arrays. Contains example trajectories of the RNN state. n_inits: int specifying the number of sampled states to return. noise_scale (optional): non-negative float specifying the standard deviation of IID Gaussian noise samples added to the sampled states. Returns: initial_states: Sampled RNN states as a [n_inits x n_states] numpy array or as an LSTMStateTuple with .c and .h as [n_inits x n_states] numpy arrays (type matches than of state_traj). Raises: ValueError if noise_scale is negative. ''' if self.is_lstm: state_traj_bxtxd = tf_utils.convert_from_LSTMStateTuple(state_traj) else: state_traj_bxtxd = state_traj [n_batch, n_time, n_states] = state_traj_bxtxd.shape # Draw random samples from state trajectories states = np.zeros([n_inits, n_states]) for init_idx in range(n_inits): trial_idx = rng.randint(n_batch) time_idx = rng.randint(n_time) states[init_idx, :] = state_traj_bxtxd[trial_idx, time_idx, :] # Add IID Gaussian noise to the sampled states if noise_scale > 0.0: states += noise_scale * rng.randn(n_inits, n_states) elif noise_scale < 0.0: raise ValueError('noise_scale must be non-negative,' ' but was %f' % noise_scale) else: # noise_scale == 0 --> don't add noise pass if self.is_lstm: return tf_utils.convert_to_LSTMStateTuple(states) else: return states
def _identify_unique_fixed_points(self): '''Identifies the unique fixed points found after optimizing from all initiali_states After running, the following class variables contain the data corresponding to the unique fixed points (see find_fixed_points for detailed descriptions): unique_xstar, unique_F_xstar, unique_qstar, unique_dq, unique_n_iters Args: None. Returns: None. ''' def unique_rows(x, approx_tol): # Quick and dirty. Can update using pdist if necessary d = int(np.round(np.max([0 - np.log10(approx_tol)]))) ux, idx = np.unique(x.round(decimals=d), axis=0, return_index=True) return ux, idx self.unique_xstar, idx = unique_rows(self.xstar, self.tol_unique) self.n_unique = len(idx) self.unique_F_xstar = self.F_xstar[idx] if self.is_lstm: self.unique_states = \ tf_utils.convert_to_LSTMStateTuple(self.unique_xstar) self.unique_new_states = \ tf_utils.convert_to_LSTMStateTuple(self.unique_F_xstar) else: self.unique_states = self.unique_xstar self.unique_new_states = self.unique_F_xstar self.unique_qstar = self.qstar[idx] self.unique_dq = self.dq[idx] self.unique_n_iters = self.n_iters[idx] '''In normal operation, Jacobians haven't yet been computed, so don't
def _grab_RNN(self, initial_states, inputs): '''Creates objects for interfacing with the RNN. These objects include 1) the optimization variables (initialized to the user-specified initial_states) which will, after optimization, contain fixed points of the RNN, and 2) hooks into those optimization variables that are required for building the TF graph. Args: initial_states: Either an [n x n_states] numpy array or an LSTMStateTuple with initial_states.c and initial_states.h as [n x n_states/2] numpy arrays. These data specify the initial states of the RNN, from which the optimization will search for fixed points. The choice of type must be consistent with state type of rnn_cell. inputs: A [n x n_inputs] numpy array specifying the inputs to the RNN for this fixed point optimization. Returns: x: An [n x n_states] tf.Variable (the optimization variable) representing RNN states, initialized to the values in initial_states. If the RNN is an LSTM, n_states represents the concatenated hidden and cell states. F: An [n x n_states] tf op representing the state transition function of the RNN applied to x. ''' if self.is_lstm: c_h_init = tf_utils.convert_from_LSTMStateTuple(initial_states) x = tf.Variable(c_h_init, dtype=self.tf_dtype) x_rnncell = tf_utils.convert_to_LSTMStateTuple(x) else: x = tf.Variable(initial_states, dtype=self.tf_dtype) x_rnncell = x n = x.shape[0] inputs_tf = tf.constant(inputs, dtype=self.tf_dtype) output, F_rnncell = self.rnn_cell(inputs_tf, x_rnncell) if self.is_lstm: F = tf_utils.convert_from_LSTMStateTuple(F_rnncell) else: F = F_rnncell init = tf.variables_initializer(var_list=[x]) self.session.run(init) return x, F
def _get_rnncell_compatible_states(self, states): '''Converts RNN states if necessary to be compatible with self.rnn_cell. Args: states: Either a numpy array or LSTMStateTuple. Returns: A representation of states that is compatible with self.rnn_cell. If self.rnn_cell is an LSTMCell, the representation is as an LSTMStateTuple. Otherwise, the representation is a numpy array. ''' if self.is_lstm: return tf_utils.convert_to_LSTMStateTuple(states) else: return states
def _run_additional_iterations_on_outliers(self): '''Detects outlier states with respect to the q function and runs additional optimization iterations on those states, using the (slow) sequential optimization procedure. This should only be used after calling either _run_joint_optimization or _run_sequential_optimizations. Updates class variables containing optimization results. Args: None. Returns: None. ''' outlier_min_q = np.median(self.qstar) * self.outlier_q_scale is_outlier = self.qstar > outlier_min_q idx_outliers = np.where(is_outlier)[0] n_outliers = len(idx_outliers) print('Detected %d \"outliers.\"' % n_outliers) if n_outliers == 0: return print('Performing additional optimization iterations.') for counter, idx in enumerate(idx_outliers): print('\n\tOutlier %d of %d (q=%.3e).' % (counter + 1, n_outliers, self.qstar[idx])) initial_state = np.expand_dims(self.xstar[idx], axis=0) if self.is_lstm: initial_state = tf_utils.convert_to_LSTMStateTuple( initial_state) xstar, F_xstar, qstar, dq, n_iters = self._run_single_optimization( initial_state) self.xstar[idx] = xstar self.F_xstar[idx] = F_xstar self.qstar[idx] = qstar self.dq[idx] = dq self.n_iters[idx] += n_iters