ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf_agents.distributions.utils.parameters_to_dict

Converts value to a nested dict (excluding all type_ info).

Sub-dicts represent Params objects; keys represent flattened nest structures in value.params.

Example:

scale_matrix = tf.Variable([[1.0, 2.0], [-1.0, 0.0]])
d = tfp.distributions.MultivariateNormalDiag(
    loc=[1.0, 1.0], scale_diag=[2.0, 3.0], validate_args=True)
b = tfp.bijectors.ScaleMatvecLinearOperator(
    scale=tf.linalg.LinearOperatorFullMatrix(matrix=scale_matrix),
    adjoint=True)
b_d = b(d)
p = utils.get_parameters(b_d)
params_dict = utils.parameters_to_dict(p, tensors_only)

results in the nested dictionary, if tensors_only=False:

{
  "bijector": {"adjoint": True,
               "scale": {"matrix": scale_matrix} },
  "distribution": {"validate_args": True,
                   # These are deeply nested because we passed lists
                   # intead of numpy arrays for `loc` and `scale_diag`.
                   "scale_diag:0": 2.0,
                   "scale_diag:1": 3.0,
                   "loc:0": 1.0,
                   "loc:1": 1.0}
}

results if tensors_only=True:

{
  "bijector": {"scale": {"matrix": scale_matrix} },
  "distribution": {},
}

The dictionary may then be modified or updated (e.g., in place), and converted back to a Params object using merge_to_parameters_from_dict.

value The (possibly recursively defined) Params.
tensors_only Whether to include all parameters or only tensors and TypeSpecs. If True, then only tf.Tensor and tf.TypeSpec leaf parameters are emitted.

A dict mapping value.params to flattened key/value pairs. Any sub-Params objects become nested dicts.