Merge pull request #3789 from verilog-to-routing/move_generator_sync_state

Add MoveGenerator methods for syncing RL agent state

Author: soheilshahrouz

vpr/src/place/move_generators/move_generator.h

@@ -146,6 +146,26 @@ class MoveGenerator {
      */
     virtual void process_outcome(double /*reward*/, e_reward_function /*reward_fun*/) {}
 
+    /**
+     * @brief Returns a token describing the generator's state after the most
+     * recent propose_move(). Only needed when outcomes are not reported right
+     * after each proposal. A caller can record the token and reward of each
+     * proposal and apply them later in a batch. Stateless generators return 0.
+     */
+    virtual size_t save_proposal_state() const { return 0; }
+
+    /**
+     * @brief Restores a token from save_proposal_state() so the next
+     * process_outcome() is credited to that proposal.
+     */
+    virtual void restore_proposal_state(size_t /*state*/) {}
+
+    /**
+     * @brief Copies the state from `other`, which must be of the same
+     * concrete type, so this generator proposes what `other` would.
+     */
+    virtual void copy_state_from(const MoveGenerator& /*other*/) {}
+
     /**
      * @brief Calculates the agent's reward and the total process outcome
      *

vpr/src/place/move_generators/simpleRL_move_generator.cpp

@@ -31,6 +31,14 @@ void SimpleRLMoveGenerator::process_outcome(double reward, e_reward_function rew
     karmed_bandit_agent->process_outcome(reward, reward_fun);
 }
 
+void SimpleRLMoveGenerator::copy_state_from(const MoveGenerator& other) {
+    // Callers always pass a generator of the same concrete type,
+    // and this runs once per sync rather than per move, so the cost is negligible.
+    const SimpleRLMoveGenerator* other_rl = dynamic_cast<const SimpleRLMoveGenerator*>(&other);
+    VTR_ASSERT_MSG(other_rl != nullptr, "Can only copy agent state from another SimpleRLMoveGenerator.");
+    karmed_bandit_agent->copy_state_from(*other_rl->karmed_bandit_agent);
+}
+
 /*                                        *
  *                                        *
  *  K-Armed bandit agent implementation   *
@@ -164,6 +172,15 @@ void KArmedBanditAgent::write_agent_info(int last_action, double reward) {
     fflush(agent_info_file_);
 }
 
+void KArmedBanditAgent::copy_state_from(const KArmedBanditAgent& other) {
+    VTR_ASSERT_SAFE(num_available_actions_ == other.num_available_actions_);
+    VTR_ASSERT_SAFE(q_.size() == other.q_.size());
+
+    exp_alpha_ = other.exp_alpha_;
+    q_ = other.q_;
+    num_action_chosen_ = other.num_action_chosen_;
+}
+
 void KArmedBanditAgent::set_step(float gamma, int move_lim) {
     if (gamma < 0) {
         exp_alpha_ = -1; //Use sample average

vpr/src/place/move_generators/simpleRL_move_generator.h

@@ -55,6 +55,17 @@ class KArmedBanditAgent {
      */
     void set_step(float gamma, int move_lim);
 
+    /// @brief Returns the action (arm) selected by the most recent propose_action() call.
+    size_t last_action() const { return last_action_; }
+
+    /// @brief Overrides the action credited by the next process_outcome() call.
+    ///        See MoveGenerator::save_proposal_state() for the intended usage.
+    void set_last_action(size_t action) { last_action_ = action; }
+
+    /// @brief Copies the learned state (Q-values, action counts, step size) from `other`.
+    ///        Both agents must have been constructed with identical configurations.
+    void copy_state_from(const KArmedBanditAgent& other);
+
   protected:
     /**
      * @brief Converts an action index to a move type.
@@ -252,6 +263,15 @@ class SimpleRLMoveGenerator : public MoveGenerator {
 
     // Receives feedback about the outcome of the previously proposed move
     void process_outcome(double reward, e_reward_function reward_fun) override;
+
+    /// @brief Saves/restores the agent action behind the most recent proposal.
+    ///        See MoveGenerator::save_proposal_state() for the intended usage.
+    size_t save_proposal_state() const override { return karmed_bandit_agent->last_action(); }
+    void restore_proposal_state(size_t state) override { karmed_bandit_agent->set_last_action(state); }
+
+    /// @brief Copies the agent state from another SimpleRLMoveGenerator.
+    ///        See MoveGenerator::copy_state_from() for the intended usage.
+    void copy_state_from(const MoveGenerator& other) override;
 };
 
 template<class T, class>