def setUp(self):
super().setUp()
self.key = jax.random.PRNGKey(0)
self.action_shape = (3, )
self.num_layers = 3
self.example_action = jnp.zeros(self.action_shape, dtype=jnp.float32)
self.example_state = jnp.zeros((8, ), dtype=jnp.float32)
self.network_def = continuous_networks.SACNetwork(
self.action_shape, self.num_layers)
self.params = self.network_def.init(self.key, self.example_state,
self.key)
def setup(self):
self._sac_network = continuous_networks.SACNetwork(
self.action_shape, self.num_layers,
self.hidden_units, self.action_limits)
def setup(self):
self._encoder = internal_continuous_networks.SACEncoderNetwork()
self._sac_network = continuous_networks.SACNetwork(
self.action_shape, self.num_layers, self.hidden_units,
self.action_limits)