def test_apply_round_robin(self): x = jnp.ones(10) # Base input bx = jnp.ones((7, 10)) # Batched input wrapped_ffn = params_adding_ffn(x) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_round_robin, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) self.assertTupleEqual(params.shape, (3, ) + x.shape) y = rr_ensemble.apply(params, jnp.broadcast_to(x, (3, ) + x.shape)) self.assertTupleEqual(y.shape, (3, ) + x.shape) np.testing.assert_allclose(params, y - x) # Note: the ensemble dimension must lead, the batch dimension is no longer # the leading dimension. by = rr_ensemble.apply( params, jnp.broadcast_to(jnp.expand_dims(bx, axis=0), (3, ) + bx.shape)) self.assertTupleEqual(by.shape, (3, ) + bx.shape) # If num_networks=3, then `round_robin(params, input)[4]` should be equal # to `apply(params[1], input[4])`, etc. yy = rr_ensemble.apply(params, jnp.broadcast_to(x, (6, ) + x.shape)) self.assertTupleEqual(yy.shape, (6, ) + x.shape) np.testing.assert_allclose(jnp.concatenate([params, params], axis=0), yy - jnp.expand_dims(x, axis=0))
def test_round_robin_random(self): x = jnp.ones(10) # Base input bx = jnp.ones((9, 10)) # Batched input ffn = RandomFFN() wrapped_ffn = networks.FeedForwardNetwork(init=functools.partial( ffn.init, x=x), apply=ffn.apply) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_round_robin, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) out = rr_ensemble.apply(params, bx) # The output should be the same every 3 rows: blocks = jnp.split(out, 3, axis=0) np.testing.assert_array_equal(blocks[0], blocks[1]) np.testing.assert_array_equal(blocks[0], blocks[2]) self.assertTrue((out[0] != out[1]).any()) for i in range(9): np.testing.assert_allclose( out[i], ffn.apply(jax.tree_map(lambda p, i=i: p[i % 3], params), bx[i]), atol=1E-5, rtol=1E-5)
def test_mean_random(self): x = jnp.ones(10) bx = jnp.ones((9, 10)) ffn = RandomFFN() wrapped_ffn = networks.FeedForwardNetwork(init=functools.partial( ffn.init, x=x), apply=ffn.apply) mean_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_mean, num_networks=3) key = jax.random.PRNGKey(0) params = mean_ensemble.init(key) single_output = mean_ensemble.apply(params, x) self.assertEqual(single_output.shape, (15, )) batch_output = mean_ensemble.apply(params, bx) # Make sure all rows are equal: np.testing.assert_allclose(jnp.broadcast_to(batch_output[0], batch_output.shape), batch_output, atol=1E-5, rtol=1E-5) # Check results explicitly: all_members = jnp.concatenate([ jnp.expand_dims(ffn.apply( jax.tree_map(lambda p, i=i: p[i], params), bx), axis=0) for i in range(3) ]) batch_means = jnp.mean(all_members, axis=0) np.testing.assert_allclose(batch_output, batch_means, atol=1E-5, rtol=1E-5)
def test_apply_mean_multiargs(self): x = jnp.ones(10) # Base input wrapped_ffn = funny_args_ffn(x) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_mean, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) self.assertTupleEqual(params.shape, (3, ) + x.shape) y = rr_ensemble.apply(params, x, 2 * x, x) self.assertTupleEqual(y.shape, x.shape) np.testing.assert_allclose(jnp.mean(params, axis=0), y - 2 * x, atol=1E-5, rtol=1E-5) y = rr_ensemble.apply(params, x, bar=x, foo=2 * x) self.assertTupleEqual(y.shape, x.shape) np.testing.assert_allclose(jnp.mean(params, axis=0), y - 2 * x, atol=1E-5, rtol=1E-5)
def test_apply_round_robin_multiargs(self): x = jnp.ones(10) # Base input wrapped_ffn = funny_args_ffn(x) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_round_robin, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) self.assertTupleEqual(params.shape, (3, ) + x.shape) ex = jnp.broadcast_to(x, (3, ) + x.shape) y = rr_ensemble.apply(params, ex, 2 * ex, ex) self.assertTupleEqual(y.shape, (3, ) + x.shape) np.testing.assert_allclose(params, y - jnp.broadcast_to(2 * x, (3, ) + x.shape), atol=1E-5, rtol=1E-5) y = rr_ensemble.apply(params, ex, bar=ex, foo=2 * ex) self.assertTupleEqual(y.shape, (3, ) + x.shape) np.testing.assert_allclose(params, y - jnp.broadcast_to(2 * x, (3, ) + x.shape), atol=1E-5, rtol=1E-5)
def make_ensemble_regressor_learner( name: str, num_networks: int, logger_fn: loggers.LoggerFactory, counter: counting.Counter, rng_key: jnp.ndarray, iterator: Iterator[types.Transition], base_network: networks_lib.FeedForwardNetwork, loss: mbop_losses.TransitionLoss, optimizer: optax.GradientTransformation, num_sgd_steps_per_step: int, ): """Creates an ensemble regressor learner from the base network. Args: name: Name of the learner used for logging and counters. num_networks: Number of networks in the ensemble. logger_fn: Constructs a logger for a label. counter: Parent counter object. rng_key: Random key. iterator: An iterator of time-batched transitions used to train the networks. base_network: Base network for the ensemble. loss: Training loss to use. optimizer: Optax optimizer. num_sgd_steps_per_step: Number of gradient updates per step. Returns: An ensemble regressor learner. """ mbop_ensemble = ensemble.make_ensemble(base_network, ensemble.apply_all, num_networks) local_counter = counting.Counter(parent=counter, prefix=name) local_logger = logger_fn(name, local_counter.get_steps_key()) if logger_fn else None def loss_fn(apply_fn: Callable[..., networks_lib.NetworkOutput], params: networks_lib.Params, key: jnp.ndarray, transitions: types.Transition) -> jnp.ndarray: del key return loss(functools.partial(apply_fn, params), transitions) # This is effectively a regressor learner. return bc.BCLearner( mbop_ensemble, rng_key, loss_fn, optimizer, iterator, num_sgd_steps_per_step, logger=local_logger, counter=local_counter)
def test_ensemble_init(self): x = jnp.ones(10) # Base input wrapped_ffn = params_adding_ffn(x) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_round_robin, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) self.assertTupleEqual(params.shape, (3, ) + x.shape) # The ensemble dimension is the lead dimension. self.assertFalse((params[0, ...] == params[1, ...]).all())
def test_apply_all_structured(self): x = jnp.ones(10) sx = [(3 * x, 2 * x), 5 * x] # Base input wrapped_ffn = struct_params_adding_ffn(sx) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_all, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) y = rr_ensemble.apply(params, sx) ex = jnp.broadcast_to(x, (3, ) + x.shape) np.testing.assert_allclose(y[0][0], params[0][0] + 3 * ex)
def test_apply_mean_structured(self): x = jnp.ones(10) sx = [(3 * x, 2 * x), 5 * x] # Base input wrapped_ffn = struct_params_adding_ffn(sx) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_mean, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) y = rr_ensemble.apply(params, sx) np.testing.assert_allclose(y[0][0], jnp.mean(params[0][0], axis=0) + 3 * x, atol=1E-5, rtol=1E-5)
def test_apply_all(self): x = jnp.ones(10) # Base input bx = jnp.ones((7, 10)) # Batched input wrapped_ffn = params_adding_ffn(x) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_all, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) self.assertTupleEqual(params.shape, (3, ) + x.shape) y = rr_ensemble.apply(params, x) self.assertTupleEqual(y.shape, (3, ) + x.shape) np.testing.assert_allclose(params, y - jnp.broadcast_to(x, (3, ) + x.shape)) by = rr_ensemble.apply(params, bx) # Note: the batch dimension is no longer the leading dimension. self.assertTupleEqual(by.shape, (3, ) + bx.shape)
def test_apply_mean(self): x = jnp.ones(10) # Base input bx = jnp.ones((7, 10)) # Batched input wrapped_ffn = params_adding_ffn(x) rr_ensemble = ensemble.make_ensemble(wrapped_ffn, ensemble.apply_mean, num_networks=3) key = jax.random.PRNGKey(0) params = rr_ensemble.init(key) self.assertTupleEqual(params.shape, (3, ) + x.shape) self.assertFalse((params[0, ...] == params[1, ...]).all()) y = rr_ensemble.apply(params, x) self.assertTupleEqual(y.shape, x.shape) np.testing.assert_allclose(jnp.mean(params, axis=0), y - x, atol=1E-5, rtol=1E-5) by = rr_ensemble.apply(params, bx) self.assertTupleEqual(by.shape, bx.shape)