Add helper function to get bargaining action by quantities.

PiperOrigin-RevId: 473450346 Change-Id: I1a963aae23ab73d22f81cf4b68bd791d120f83c3
google-deepmind · Sep 11, 2022 · fec7620 · fec7620
1 parent 7fa9941
commit fec7620
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/open_spiel/games/bargaining.cc b/open_spiel/games/bargaining.cc
@@ -482,6 +482,17 @@ int BargainingGame::NumDistinctActions() const {
   return all_offers_.size() + 1;
 }
 
+std::pair<Offer, Action> BargainingGame::GetOfferByQuantities(
+    const std::vector<int>& quantities) const {
+  for (int i = 0; i < all_offers_.size(); ++i) {
+    if (quantities == all_offers_[i].quantities) {
+      return {all_offers_[i], i};
+    }
+  }
+  return {Offer(), kInvalidAction};
+}
+
+
 std::vector<int> BargainingGame::ObservationTensorShape() const {
   return {
       1 +                                       // Agreement reached?

diff --git a/open_spiel/games/bargaining.h b/open_spiel/games/bargaining.h
@@ -157,6 +157,8 @@ class BargainingGame : public Game {
   const std::vector<Offer>& AllOffers() const { return all_offers_; }
   const Instance& GetInstance(int num) const { return all_instances_[num]; }
   const Offer& GetOffer(int num) const { return all_offers_[num]; }
+  std::pair<Offer, Action> GetOfferByQuantities(
+      const std::vector<int>& quantities) const;
 
  private:
   void ParseInstancesFile(const std::string& filename);

diff --git a/open_spiel/python/pybind11/games_bargaining.cc b/open_spiel/python/pybind11/games_bargaining.cc
@@ -24,6 +24,7 @@ using open_spiel::State;
 using open_spiel::bargaining::BargainingGame;
 using open_spiel::bargaining::BargainingState;
 using open_spiel::bargaining::Instance;
+using open_spiel::bargaining::Offer;
 
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(BargainingGame);
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(BargainingState);
@@ -34,6 +35,10 @@ void open_spiel::init_pyspiel_games_bargaining(py::module& m) {
       .def_readwrite("pool", &Instance::pool)
       .def_readwrite("values", &Instance::values);
 
+  py::class_<Offer>(m, "Offer")
+      .def(py::init<>())
+      .def_readwrite("quantities", &Offer::quantities);
+
   py::classh<BargainingState, State>(m, "BargainingState")
       .def("instance", &BargainingState::instance)
       .def("agree_action", &BargainingState::AgreeAction)
@@ -54,6 +59,9 @@ void open_spiel::init_pyspiel_games_bargaining(py::module& m) {
 
   py::classh<BargainingGame, Game>(m, "BargainingGame")
     .def("all_instances", &BargainingGame::AllInstances)
+    // get_offer_by_quantities(quantities: List[int]). Returns a tuple
+    // of (offer, OpenSpiel action)
+    .def("get_offer_by_quantities", &BargainingGame::GetOfferByQuantities)
     // Pickle support
     .def(py::pickle(
         [](std::shared_ptr<const BargainingGame> game) {  // __getstate__