code: def prepare_for_train(self):
      self.step_count = 0
  
      self.add_node_agent.prepare_training()
      self.add_edge_agent.prepare_training()
      self.choose_dest_agent.prepare_training()
  
imports: from functools import partial
  import dgl
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.distributions import Bernoulli, Categorical
label_desc: Use image node_id 4 for calling the DGMG obj's underlying member method code with example usage: obj.prepare_for_train() without return types