tf_agents.utils.common.function_in_tf1

View source on GitHub

Wrapper that returns common.function if using TF1.

This allows for code that assumes autodeps is available to be written once, in the same way, for both TF1 and TF2.

Usage:

train = function_in_tf1()(agent.train)
loss = train(experience)

*args Arguments for common.function.
**kwargs Keyword arguments for common.function.

A callable that wraps a function.