def __init__(self, inputs, model, input_lens): # skip TFActionDistribution init ActionDistribution.__init__(self, inputs, model) self.cats = [ Categorical(input_, model) for input_ in tf.split(inputs, input_lens, axis=1) ] self.sample_op = self._build_sample_op()
def __init__(self, inputs, model, action_space, child_distributions, input_lens): # skip TFActionDistribution init ActionDistribution.__init__(self, inputs, model) self.input_lens = input_lens split_inputs = tf.split(inputs, self.input_lens, axis=1) child_list = [] for i, distribution in enumerate(child_distributions): child_list.append(distribution(split_inputs[i], model)) self.child_distributions = child_list
def __init__(self, inputs, model): # skip TFActionDistribution init ActionDistribution.__init__(self, inputs, model) input_lens = inputs.get_shape().as_list()[1] self.gausses = [ SquashedGaussian(tf.squeeze(input_, axis=1), model) for input_ in tf.split(inputs, input_lens, axis=1) ] self.sample_op = self._build_sample_op() self.sampled_action_logp_op = self.logp(self.sample_op)
def __init__(self, inputs, model): # skip TFActionDistribution init ActionDistribution.__init__(self, inputs, model) input_lens = inputs.get_shape().as_list()[1] self.cats = [ Categorical(input_, model) for input_ in tf.unstack(inputs, input_lens, axis=1) ] self.sample_op = self._build_sample_op() self.sampled_action_logp_op = self.logp(self.sample_op)
def __init__(self, inputs: List[TensorType], model: ModelV2, input_lens: Union[List[int], np.ndarray, Tuple[int, ...]]): # skip TFActionDistribution init ActionDistribution.__init__(self, inputs, model) self.cats = [ Categorical(input_, model) for input_ in tf.split(inputs, input_lens, axis=1) ] self.sample_op = self._build_sample_op() self.sampled_action_logp_op = self.logp(self.sample_op)
def __init__(self, inputs, model, *, child_distributions, input_lens, action_space): ActionDistribution.__init__(self, inputs, model) self.action_space_struct = get_base_struct_from_space(action_space) self.input_lens = np.array(input_lens, dtype=np.int32) split_inputs = tf.split(inputs, self.input_lens, axis=1) self.flat_child_distributions = tree.map_structure( lambda dist, input_: dist(input_, model), child_distributions, split_inputs )
def __init__(self, inputs: List[TensorType], model: ModelV2, input_lens: Union[List[int], np.ndarray, Tuple[int, ...]], action_space=None): # skip TFActionDistribution init ActionDistribution.__init__(self, inputs, model) self.cats = [ Categorical(input_, model) for input_ in tf.split(inputs, input_lens, axis=1) ] self.action_space = action_space if self.action_space is None: self.action_space = gym.spaces.MultiDiscrete( [c.inputs.shape[1] for c in self.cats]) self.sample_op = self._build_sample_op() self.sampled_action_logp_op = self.logp(self.sample_op)