Demo: Knapsack Problem¶

# knapsack problem setup
num_items = 5
max_capacity = 20
key = jr.PRNGKey(0)
key_w, key_r = jr.split(key)
item_weights = max_capacity / num_items + jr.uniform(key_w, shape=(num_items,), minval=0.0, maxval=5.0)
rewards = jr.uniform(key_r, shape=(num_items,), minval=0.0, maxval=1.0)

print("item rewards", rewards)
print("item weights", item_weights)
print("item weight sum", item_weights.sum())

# mdp config
state_config = {
    f"s_{i}": {"elements": ["not enclude", "include"], "depends_on": [f"s_{i}"], "controlled_by": [f"a_{i}"]} 
    for i in range(num_items)
}
state_config["z"] = {
    "elements": ["not violated", "violated"], "depends_on": ["z"], "controlled_by": [f"a_{i}" for i in range(num_items)]
} 

obs_config = {
    k: {"elements": v["elements"], "depends_on": [k]} for k, v in state_config.items()
}

act_config = {
    f"a_{i}": {"elements": ["not enclude", "include"]} 
    for i in range(num_items)
}

model_description = {
    "observations": obs_config,
    "controls": act_config,
    "states": state_config,
}

model = distribution.compile_model(model_description)

print("A shapes", [a.data.shape for a in model.A])
print("B shapes", [a.data.shape for a in model.B])

def create_identity_transition_factor(mat):
    for i in range(mat.shape[1]):
        mat[:, i] = np.eye(len(mat))
    return mat

def create_constraint_factor_z_greater_than(act_dim, maximum, num_items, weights):
    # Create an array of shape (2, act_dim, act_dim, ..., act_dim)
    tensor_shape = (2,) + (act_dim,) * num_items
    
    # Create an array with indices from 0 to act_dim - 1 along each dimension
    indices = np.indices(tensor_shape[1:])

    # Reshape weights to fit indices shape
    weights_reshaped = np.array(weights).reshape((-1,) + (1,) * (indices.ndim - 1))
    # Multiply weights with matrix that conforms to constraint
    result = np.array(indices == (act_dim - 1)) * weights_reshaped

    # Calculate the total for each combination of actions
    total = np.sum(result, axis=0)
    
    # Create the tensor based on the total hours condition
    tensor = np.where(total > maximum, 1, 0)

    # Stack the tensor along the first axis to create the final tensor
    tensor = np.stack((1 - tensor, tensor), axis=0)

    # make a copy for self state 
    tensor = np.stack([tensor, tensor], axis=1)
    return tensor

# update A tensor
for i in range(len(model.A)):
    model.A[i].data = np.eye(len(model.A[i].data))

# update B tensors
for i in range(num_items):
    model.B[i].data = create_identity_transition_factor(model.B[i].data)

model.B[-1].data = create_constraint_factor_z_greater_than(2, max_capacity, num_items, item_weights)

# create C tensors
preferences = nn.softmax(np.stack([np.zeros_like(rewards), rewards], axis=-1), axis=-1)
Cs = [None for _ in range(len(model.A))]
for i in range(len(model.A)):
    Cs[i] = preferences[i]
    
Cs[-1] = np.array([1., 0]) # capacity constraint cannot be violated

print("A shapes", [a.data.shape for a in model.A])
print("B shapes", [a.data.shape for a in model.B])

print("A normalized", [np.isclose(a.data.sum(0), 1.).all() for a in model.A])
print("B normalized", [np.isclose(b.data.sum(0), 1.).all() for b in model.B])
print("C normalized", [np.isclose(a.sum(0), 1.).all() for a in Cs])

print("best action", action)
print("best action multi", action_multi)
print("item weights")
print(item_weights)
print("item rewards")
print(rewards)

def brute_force_knapsack(weights, rewards, capacity):
    best_r = -jnp.inf
    best_bits = None
    for bits in itertools.product([0,1], repeat=len(weights)):
        w = jnp.dot(jnp.array(bits), weights)
        if w <= capacity + 1e-9:
            r = jnp.dot(jnp.array(bits), rewards)
            if r > best_r:
                best_r, best_bits = float(r), jnp.array(bits)
    return best_bits, best_r

bf_bits, bf_reward = brute_force_knapsack(item_weights, rewards, max_capacity)
print("Brute-force best bits:", bf_bits, "reward:", bf_reward)
print("Agent best bits:", action_multi[0], "reward:", float(jnp.dot(action_multi[0], rewards)))

# print top actions, sorted in decreasing posterior probability (lowest EFE first)
from pymdp import utils
for i, idx in enumerate(np.argsort(q_pi[0])[::-1]):
    action_multi_f = utils.index_to_combination(agent.policies[idx, 0][-1].tolist(), agent.num_controls_multi)
    print("\naction", agent.policies[idx, 0])
    print("action multi", action_multi_f)
    print("efe: -{:.2f}, reward: {:.2f}".format(G[0, idx], np.sum(rewards * action_multi_f)))
    if i == 5:
        break