test: [RLOS_2023] test for contextual bandit

python/tests/assert_job.py

@@ -1,18 +1,24 @@
 import numpy as np
-from numpy.testing import assert_allclose, assert_array_almost_equal
+from numpy.testing import assert_allclose, assert_almost_equal
 from vw_executor.vw import ExecutionStatus
 
 
+def remove_non_digits(string):
+    return "".join(char for char in string if char.isdigit() or char == ".")
+
+
 def get_from_kwargs(kwargs, key, default=None):
     if key in kwargs:
         return kwargs[key]
     else:
         return default
 
+
 def majority_close(arr1, arr2, rtol, atol, threshold):
     # Check if the majority of elements are close
     close_count = np.count_nonzero(np.isclose(arr1, arr2, rtol=rtol, atol=atol))
-    return close_count > len(arr1) * threshold
+    return close_count >= len(arr1) * threshold
+
 
 def assert_weight(job, **kwargs):
     atol = get_from_kwargs(kwargs, "atol", 10e-8)
@@ -23,23 +29,39 @@ def assert_weight(job, **kwargs):
     with open(data[0], "r") as f:
         data = f.readlines()
     data = [i.strip() for i in data]
-    weights = job[0].model9('--readable_model').weights
+    weights = job[0].model9("--readable_model").weights
     weights = weights["weight"].to_list()
-    assert_allclose(weights, expected_weights, atol=atol, rtol=rtol), f"weights should be {expected_weights}"
+    assert_allclose(
+        weights, expected_weights, atol=atol, rtol=rtol
+    ), f"weights should be {expected_weights}"
 
-def assert_prediction(job, **kwargs):
-        assert job.status == ExecutionStatus.Success, "job should be successful"
-        atol = kwargs.get("atol", 10e-8)
-        rtol = kwargs.get("rtol", 10e-5)
-        threshold = kwargs.get("threshold", 0.9)
-        constant = kwargs["expected_value"]
-        predictions = job.outputs['-p']
-        with open(predictions[0], "r") as f:
-            predictions = f.readlines()
-            predictions = [float(i) for i in predictions[1:]]
-            assert majority_close(predictions, [constant]*len(predictions), rtol=rtol, atol=atol, threshold=threshold), f"predicted value should be {constant}"
 
+def assert_prediction(job, **kwargs):
+    assert job.status == ExecutionStatus.Success, "job should be successful"
+    atol = kwargs.get("atol", 10e-8)
+    rtol = kwargs.get("rtol", 10e-5)
+    threshold = kwargs.get("threshold", 0.9)
+    expected_value = kwargs["expected_value"]
+    predictions = job.outputs["-p"]
+    with open(predictions[0], "r") as f:
+        prediction = [i.strip() for i in f.readlines()]
+        prediction = [i for i in prediction if i != ""]
+        if ":" in prediction[0]:
+            prediction = [[j.split(":")[1] for j in i.split(",")] for i in prediction]
+        if type(prediction[0]) == list:
+            prediction = [[float(remove_non_digits(j)) for j in i] for i in prediction]
+        else:
+            prediction = [float(remove_non_digits(i)) for i in prediction]
+        assert majority_close(
+            prediction,
+            [expected_value] * len(prediction),
+            rtol=rtol,
+            atol=atol,
+            threshold=threshold,
+        ), f"predicted value should be {expected_value}, \n actual values are {prediction}"
 
 
-def assert_functions():
-    return 
+def assert_loss(job, **kwargs):
+    assert job.status == ExecutionStatus.Success, "job should be successful"
+    assert type(job[0].loss) == float, "loss should be an float"
+    assert_almost_equal(job[0].loss, kwargs["expected_loss"], decimal=2)

python/tests/data_generation.py

@@ -1,10 +1,75 @@
 import random
+import os
+from test_helper import get_function_object
+
+script_directory = os.path.dirname(os.path.realpath(__file__))
+random.seed(10)
+
 
 def constant_function(no_sample, constant, lower_bound, upper_bound):
     dataFile = f"constant_func_{no_sample}_{constant}_{upper_bound}_{lower_bound}.txt"
-    with open(dataFile, "w") as f:
-        random.seed(10)
+    with open(os.path.join(script_directory, dataFile), "w") as f:
         for _ in range(no_sample):
             x = random.uniform(lower_bound, upper_bound)
             f.write(f"{constant} |f x:{x}\n")
     return dataFile
+
+
+def random_number_items(items):
+    num_items_to_select = random.randint(1, len(items))
+    return random.sample(items, num_items_to_select)
+
+
+def generate_cb_data(
+    num_examples,
+    num_features,
+    num_actions,
+    reward_function,
+    logging_policy,
+    no_context=1,
+    context_name=None,
+):
+
+    dataFile = f"cb_test_{num_examples}_{num_actions}_{num_features}.txt"
+
+    reward_function_obj = get_function_object(
+        "reward_functions", reward_function["name"]
+    )
+    logging_policy_obj = get_function_object("logging_policies", logging_policy["name"])
+    features = [f"feature{index}" for index in range(1, num_features + 1)]
+    with open(os.path.join(script_directory, dataFile), "w") as f:
+        for _ in range(num_examples):
+            if no_context > 1:
+                context = random.randint(1, no_context)
+                if not context_name:
+                    context_name = [f"{index}" for index in range(1, no_context + 1)]
+
+            def return_cost_probability(chosen_action, context=1):
+                cost = reward_function_obj(
+                    chosen_action, context, **reward_function["params"]
+                )
+                logging_policy["params"]["chosen_action"] = chosen_action
+                probability = logging_policy_obj(**logging_policy["params"])
+                return cost, probability
+
+            chosen_action = random.randint(1, num_actions)
+            if no_context > 1:
+                f.write(f"shared | User s_{context_name[context-1]}\n")
+                for action in range(1, num_actions + 1):
+
+                    cost, probability = return_cost_probability(action, context)
+                    if action == chosen_action:
+                        f.write(
+                            f'{action}:{cost}:{probability} | {" ".join(random_number_items(features))}\n'
+                        )
+                    else:
+                        f.write(f'| {" ".join(random_number_items(features))}\n')
+
+            else:
+
+                cost, probability = return_cost_probability(chosen_action)
+                f.write(
+                    f'{chosen_action}:{cost}:{probability} | {" ".join(random_number_items(features))}\n'
+                )
+            f.write("\n")
+    return dataFile

python/tests/logging_policies.py

@@ -0,0 +1,7 @@
+def constant_probability(chosen_action, **kwargs):
+    return 1
+
+
+def even_probability(chosen_action, **kwargs):
+    num_actions = kwargs["num_actions"]
+    return round(1 / num_actions, 2)

python/tests/reward_functions.py

@@ -0,0 +1,21 @@
+
+
+def fixed_reward(chosen_action, context, **kwargs):
+    return 1
+
+
+def constant_reward(chosen_action, context, **kwargs):
+    reward = kwargs["reward"]
+    return reward[chosen_action - 1]
+
+
+def fixed_reward_two_action(chosen_action, context, **kwargs):
+    if context == 1 and chosen_action == 2:
+        return 1
+    elif context == 2 and chosen_action == 2:
+        return 0
+    elif context == 1 and chosen_action == 1:
+        return 0
+    elif context == 2 and chosen_action == 1:
+        return 1
+    return 1