Errors when saving ranking policies

[tottenjordan]

What is recommended way to save ranking policies?

For each agent policy in the ranking tutorial, im getting errors like below when trying to save with PolicySaver

saver = policy_saver.PolicySaver(
    policy = agent.policy, 
    batch_size=None
)

penalize_cosine_distance_ranking_policy error:

ValueError: Could not determine an appropriate value for field `logits` in object `tfp.distributions.CosinePenalizedPlackettLuce("PlackettLuce", batch_shape=[?], 
event_shape=[3], dtype=int32)`. Looked for 
 1. an attr called `logits`,
 2. an attr called `_logits`,
 3. an entry in `obj.parameters` with key "logits".

no_penalty_ranking_policy error:

ValueError: Could not determine an appropriate value for field `features` in object `tfp.distributions.NoPenaltyPlackettLuce("PlackettLuce", batch_shape=[?], event_shape=[50], dtype=int32)`. 
Looked for 
 1. an attr called `features`,
 2. an attr called `_features`,
 3. an entry in `obj.parameters` with key "features".

descending_score_ranking_policy error:

TypeError: To be compatible with tf.function, Python functions must return zero or more Tensors or ExtensionTypes or None values; in compilation of <function PolicySaver.__init__.<locals>.polymorphic_distribution_fn at 0x7ced7451f010>, found return value of type DescendingScoreSampler, which is not a Tensor or ExtensionType.

[SimonEOA]

Hi @tottenjordan I'm having a similar issue, did you ever get an answer, or find a solution?