diff --git a/miplearn/problems/knapsack.py b/miplearn/problems/knapsack.py
index ec4e182..a98fd0d 100644
--- a/miplearn/problems/knapsack.py
+++ b/miplearn/problems/knapsack.py
@@ -19,7 +19,7 @@ class ChallengeA:
     """
     def __init__(self,
                  seed=42,
-                 n_training_instances=300,
+                 n_training_instances=500,
                  n_test_instances=50):
 
         np.random.seed(seed)
@@ -81,16 +81,18 @@ class MultiKnapsackInstance(Instance):
         
     def get_instance_features(self):
         return np.hstack([
-            self.prices,
+            np.mean(self.prices),
             self.capacities,
-            self.weights.ravel(),
         ])
 
     def get_variable_features(self, var, index):
-        return np.array([])
+        return np.hstack([
+            self.prices[index],
+            self.weights[:, index],
+        ])
 
-    def get_variable_category(self, var, index):
-        return index
+#     def get_variable_category(self, var, index):
+#         return index
 
 
 class MultiKnapsackGenerator:
@@ -212,4 +214,34 @@ class MultiKnapsackGenerator:
             return MultiKnapsackInstance(p, b, w)
         return [_sample() for _ in range(n_samples)]
     
-    
\ No newline at end of file
+
+class KnapsackInstance(Instance):
+    """
+    Simpler (one-dimensional) Knapsack Problem, used for testing.
+    """
+    def __init__(self, weights, prices, capacity):
+        self.weights = weights
+        self.prices = prices
+        self.capacity = capacity
+
+    def to_model(self):
+        model = pe.ConcreteModel()
+        items = range(len(self.weights))
+        model.x = pe.Var(items, domain=pe.Binary)
+        model.OBJ = pe.Objective(rule=lambda m: sum(m.x[v] * self.prices[v] for v in items),
+                                 sense=pe.maximize)
+        model.eq_capacity = pe.Constraint(rule=lambda m: sum(m.x[v] * self.weights[v]
+                                                             for v in items) <= self.capacity)
+        return model
+
+    def get_instance_features(self):
+        return np.array([
+            self.capacity,
+            np.average(self.weights),
+        ])
+
+    def get_variable_features(self, var, index):
+        return np.array([
+            self.weights[index],
+            self.prices[index],
+        ])
\ No newline at end of file
diff --git a/miplearn/problems/tests/test_knapsack.py b/miplearn/problems/tests/test_knapsack.py
index 31d31da..2d6b1b6 100644
--- a/miplearn/problems/tests/test_knapsack.py
+++ b/miplearn/problems/tests/test_knapsack.py
@@ -25,24 +25,6 @@ def test_knapsack_generator():
     assert round(np.mean(b_sum), -3) == 25000.
 
     
-def test_knapsack_instance():
-    instance = MultiKnapsackInstance(
-        prices=np.array([5.0, 10.0, 15.0]),
-        capacities=np.array([20.0, 30.0]),
-        weights=np.array([
-            [5.0,  5.0,  5.0],
-            [5.0, 10.0, 15.0],
-        ])
-    )
-    
-    assert (instance.get_instance_features() == np.array([
-        5.0, 10.0, 15.0, 20.0, 30.0, 5.0, 5.0, 5.0, 5.0, 10.0, 15.0
-    ])).all()
-    
-    solver = LearningSolver()
-    results = solver.solve(instance)
-    assert results["Problem"][0]["Lower bound"] == 30.0
-    
 def test_knapsack_fixed_weights_jitter():
     gen = MultiKnapsackGenerator(n=randint(low=50, high=51),
                                  m=randint(low=10, high=11),