Training Sequence Diagrams - m-marini/wheellyj GitHub Wiki

sequenceDiagram
  participant RobotController
  participant WorldModeller
  participant Environment
  participant InferenceMediator
  participant Agent
  participant TDNetwork

  activate RobotController
  RobotController->>+WorldModeller: onLatch(RobotStatus)
  WorldModeller->>WorldModeller: update current model
  deactivate WorldModeller

  RobotController->>+WorldModeller: onInference(RobotStatus)
  WorldModeller->>WorldModeller: update model for inference

  WorldModeller->>+Environment: onInference(WorldModel)
  Environment->>Environment: compute state

  Environment->>+InferenceMediator: act(State)

  InferenceMediator->>+Agent: act(State)
  Agent->>Agent: compute input signals

  Agent->>+TDNetwork: forward(signals)
  TDNetwork-->>-Agent: output policy

  Agent->>Agent: choose actions and compute action signals
  Agent-->>-InferenceMediator: action signals

  InferenceMediator-->>-Environment: action signals
  Environment->>Environment: compute robot command and execution results

  Environment->>+InferenceMediator: observe(ExecutionResult)

  InferenceMediator->>+Agent: observe(ExecutionResult)
  Agent->>Agent: compute trajectory
  Agent-->>-InferenceMediator: trajectory

  opt ready to train
    InferenceMediator-)+Agent: trainByTrajectory(trajectory)
  end
  InferenceMediator-->>-Environment: return

  Environment-->>-WorldModeller: robot commands

  WorldModeller->>-RobotController: execute(commands)

  deactivate RobotController

  %% parallel training
  Note over Agent: prepare data from trajectory
  Agent->>+Agent: trainBatch(states, actionMasks, rewards)
  loop for each epochs
    Note over Agent: compute actionProb0

    Agent->>+Agent: trainEpoch(states, actionMasks, rewards, actionProb0)
    loop for each mini batch
      Note over Agent: compute mini batch data

      Agent->>+Agent: trainMiniBatch(states, actionMasks, rewards, pi0)
      Agent->>+TDNetwork: train(...)
      TDNetwork-->>-Agent: TDNetwork

      deactivate Agent
    end
    deactivate Agent

  end
  deactivate Agent
  deactivate Agent

sequenceDiagram
  actor BatchTraining
  participant SignalGenerator
  participant WorldModeller
  participant Environment as WorldEnvironment
  participant State
  participant Agent as AbstractAgent

  activate BatchTraining

  BatchTraining->>+SignalGenerator: generate()
  loop while not eof
    SignalGenerator->>SignalGenerator:readData()
    note over SignalGenerator: read model and command record

    SignalGenerator->>+WorldModeller:updateForInference(model)
    WorldModeller-->>-SignalGenerator: updated model

    SignalGenerator->>+Environment:state(model)
    Environment-->>-SignalGenerator:state

    SignalGenerator->>+State:signals()
    State-->>-SignalGenerator:signals

    SignalGenerator->>+Agent: processSignals(signals)
    Agent-->>-SignalGenerator:processed signals

    SignalGenerator->>SignalGenerator:writeSignals(processedSignals)

    SignalGenerator->>+Environment:actions(command)
    Environment-->>-SignalGenerator:actions

    SignalGenerator->>SignalGenerator:writeAction(actions)

    SignalGenerator->>+Environment:rewards(state0, actions, state1)
    Environment-->>-SignalGenerator:reward

    SignalGenerator->>SignalGenerator:writeReward
  end
  SignalGenerator-->>-BatchTraining:return
  deactivate BatchTraining

sequenceDiagram
  actor BatchTraining
  participant BatchTrainer
  participant Agent
  participant TDNetwork

  BatchTraining->>+BatchTrainer:validate(File)
  Note right of BatchTrainer: validate signal, reward, action files

  BatchTraining->>BatchTrainer:prepare()
  Note right of BatchTrainer: write action mask file

  BatchTraining->>BatchTrainer:train()
  BatchTrainer->>BatchTrainer: createActionProbabilities()
  loop for each epoch
    loop for each mini batch
    Note right of BatchTrainer: prepare mini batch data by reading datasets
      BatchTrainer->>+Agent: trainMiniBatch(states, actionMasks, rewards, actionProb0)

      Agent->>+TDNetwork: forward(state0)
      Note over Agent,TDNetwork: compute differential<br/>value function prediction 
      TDNetwork-->>-Agent: network state

      Agent->>+TDNetwork: forward(state0)
      Note over Agent, TDNetwork: forward pass for training
      TDNetwork-->>-Agent: network state
    
      Agent->>Agent:computeAdvantage(rewards, avgRewards, vPrediction)
      Agent->>Agent:policy(result0)
      Agent->>Agent:optimizerGrad(pi, actionMasks, actionProb0, adv)
      Note over Agent:Computes output gradients for network (merges critic and policy grads)

      Agent->>+TDNetwork:train(...)
      TDNetwork-->>-Agent: trained network

      deactivate Agent
    end
  end
  deactivate BatchTrainer