Expose Leduc state internals and action types to Python.

PiperOrigin-RevId: 511478011 Change-Id: I8ea980f14d52050d7afd1b4433ba8bf671932fa2
willmcgowan · Feb 23, 2023 · 649a5f5 · 649a5f5
1 parent 490695b
commit 649a5f5
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
diff --git a/open_spiel/games/leduc_poker.h b/open_spiel/games/leduc_poker.h
@@ -100,6 +100,19 @@ class LeducState : public State {
   // Gets the private cards.
   std::vector<int> GetPrivateCards() const { return private_cards_; }
 
+  // Gets the public card.
+  int GetPublicCard() const { return public_card_; }
+
+  // Gets number of chips in pot.
+  int GetPot() const { return pot_; }
+
+  // Gets how much money each player has.
+  std::vector<double> GetMoney() const { return money_; }
+
+  // Gets the action sequence of rounds 1 & 2.
+  std::vector<int> GetRound1() const { return round1_sequence_; }
+  std::vector<int> GetRound2() const { return round2_sequence_; }
+
   // Sets the private cards to specific ones. Note that this function does not
   // change the history, so any functions relying on the history will not longer
   // work properly.

diff --git a/open_spiel/python/pybind11/games_leduc_poker.cc b/open_spiel/python/pybind11/games_leduc_poker.cc
@@ -22,15 +22,35 @@ namespace py = ::pybind11;
 using open_spiel::Game;
 using open_spiel::State;
 using open_spiel::leduc_poker::LeducState;
+using open_spiel::leduc_poker::ActionType;
 
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(LeducState);
 
 void open_spiel::init_pyspiel_games_leduc_poker(py::module& m) {
-  py::classh<LeducState, State>(m, "LeducState")
+  py::module_ leduc_poker = m.def_submodule("leduc_poker");
+
+  leduc_poker.attr("INVALID_CARD") = py::int_(
+    open_spiel::leduc_poker::kInvalidCard);
+
+  py::enum_<ActionType>(leduc_poker, "ActionType")
+    .value("FOLD", ActionType::kFold)
+    .value("CALL", ActionType::kCall)
+    .value("RAISE", ActionType::kRaise)
+    .export_values();
+
+  py::classh<LeducState, State>(leduc_poker, "LeducState")
       // Gets the private cards; no arguments, returns vector of ints.
       .def("get_private_cards", &LeducState::GetPrivateCards)
       // Sets the private cards; takes a vector of ints, no returns.
       .def("set_private_cards", &LeducState::SetPrivateCards)
+      // Expose additional state features.
+      .def("private_card", &LeducState::private_card)
+      .def("public_card", &LeducState::public_card)
+      .def("round", &LeducState::round)
+      .def("money", &LeducState::GetMoney)
+      .def("pot", &LeducState::GetPot)
+      .def("round1", &LeducState::GetRound1)
+      .def("round2", &LeducState::GetRound2)
       // Pickle support
       .def(py::pickle(
           [](const LeducState& state) {  // __getstate__