How to bound input dimension for differential_evolution_minimize in tensorflow-probability?

Question

Unlike in scipy's implementation of differential evolution (DE), there is no direct way to define bounds for my inputs in tensorflow-probability.

Since the inputs to my function are 5-tuples defining a pixel in an image (x,y,R,G,B), the values would need to be integers and bounded to my image dimension as well as RGB values of 0 to 255 during optimization.

To paint the bigger picture: I want to minimize the confidence of an image classifier which itself is part of a reinforcement learning agent that needs to decide on an action depending on its observation of the environment. The agent is fully trained but now I want to throw a few stones by perturbing a single pixel in its observation and monitor its performance.

I want to use the differential evolution algorithm to find the pixel which would decrease the confidence of the agent in its actions the most.

Currently I have an action-prediction function, which takes as an argument the perturbation pixel, runs the perturbed observation through the classifier and returns the resulting confidence in the action which the agent would have chosen without the perturbation:

Code block I

#random perturbation pixel as an example of my input:
pixel = tf.constant([36,48,255,255,255]) # (x,y,R,G,B)

def predict_action(pixel):
   perturbed_obs = perturb_obs(pixel, observation)
   confidence = classifier(perturbed_obs)
   return confidence

Now I would want to give this function to the optimizer with an initial population:

Code block II

popsize=80

init_pop = generate_population(popsize)
# returns Tensor("scan/while/Squeeze:0", shape=(80, 5), dtype=int64)
# i.e. 80 random perturbation pixels

results = tfp.optimizer.differential_evolution_minimize(
    predict_action, initial_population=init_pop, seed=42)

Yet, how do I define the bounds of my input, such that the population will always be valid pixels?

I asked about this on GitHub and a possible way of achieving this is through the use of their bijector functionality:

Code block III

# First we squash `pixel_logits` to (0, 1), then scale it to (0, 255).
bijector = tfb.Affine(scale=255.)(tfb.Sigmoid())

def unconstrained_objective_fn(pixel_logits):
    return objective_fn(bijector.forward(pixel_logits))

results = minimize(unconstrained_objective_fn, initial_position=bijector.inverse(initial_pixels))
pixels = bijector.forward(results.position)

While I understand the approach in principle, I fail to apply this to my current situation/understanding of my problem.

EDIT: Removed information that was unrelated to the main question. Always mind your types and dimensions!

Answer 1

Based on my experience with tfp-0.6 and tf-1.13.1, Code block III could be rewritten as follows:

width = ... #some Python float
height = ... #some other Python float

bijectors = [
    tfb.Chain([tfb.AffineScalar(scale=width), tfb.Sigmoid()]),
    tfb.Chain([tfb.AffineScalar(scale=height), tfb.Sigmoid()]),
    tfb.Chain([tfb.AffineScalar(scale=255.), tfb.Sigmoid()]),
    tfb.Chain([tfb.AffineScalar(scale=255.), tfb.Sigmoid()]),
    tfb.Chain([tfb.AffineScalar(scale=255.), tfb.Sigmoid()])
  ]

def constrained_objective_fn(pixel_logits):
    constrained_pixel_logits = [b.forward(p) for b, p in zip(bijectors, pixel_logits)]
    return objective_fn(constrained_pixel_logits)

results = tfp.optimizer.differential_evolution_minimize(
constrained_objective_fn, 
initial_population=init_pop, 
seed=42)

pixels = [b.forward(p) for b, p in zip(bijectors, results.position)]

Note that the initial population from Code Block II is preserved and bijector.inverse is not applied.