Prediction: refactor pedestrian interaction model to update iteratively

modules/prediction/evaluator/pedestrian/pedestrian_interaction_evaluator.cc

@@ -112,18 +112,18 @@ bool PedestrianInteractionEvaluator::Evaluate(Obstacle* obstacle_ptr) {
           .toTensor()
           .to(torch::kCPU);
 
+  // Step 2 Get position embedding
+  double pos_x = feature_values[2];
+  double pos_y = feature_values[3];
   double rel_x = 0.0;
   double rel_y = 0.0;
   if (obstacle_ptr->history_size() > 1) {
     rel_x = obstacle_ptr->latest_feature().position().x() -
-            obstacle_ptr->earliest_feature().position().x();
+        obstacle_ptr->feature(1).position().x();
     rel_y = obstacle_ptr->latest_feature().position().y() -
-            obstacle_ptr->earliest_feature().position().y();
+        obstacle_ptr->feature(1).position().y();
   }
 
-  // Step 2 Get position embedding
-  double pos_x = feature_values[2];
-  double pos_y = feature_values[3];
   torch::Tensor torch_position = torch::zeros({1, 2});
   torch_position[0][0] = rel_x;
   torch_position[0][1] = rel_y;
@@ -140,6 +140,7 @@ bool PedestrianInteractionEvaluator::Evaluate(Obstacle* obstacle_ptr) {
   for (int i = 0; i < kEmbeddingSize; ++i) {
     lstm_input[0][i] = position_embedding[0][i];
   }
+
   if (obstacle_id_lstm_state_map_.find(id) ==
       obstacle_id_lstm_state_map_.end()) {
     obstacle_id_lstm_state_map_[id].ht = torch::zeros({1, 1, kHiddenSize});
@@ -148,10 +149,12 @@ bool PedestrianInteractionEvaluator::Evaluate(Obstacle* obstacle_ptr) {
   }
   torch::Tensor curr_ht = obstacle_id_lstm_state_map_[id].ht;
   torch::Tensor curr_ct = obstacle_id_lstm_state_map_[id].ct;
+
   for (int i = 0; i < kHiddenSize; ++i) {
     lstm_input[0][kEmbeddingSize + i] = curr_ht[0][0][i];
     lstm_input[0][kEmbeddingSize + kHiddenSize + i] = curr_ct[0][0][i];
   }
+
   std::vector<torch::jit::IValue> lstm_inputs;
   lstm_inputs.push_back(lstm_input.to(device_));
   auto lstm_out_tuple = torch_single_lstm_ptr_->forward(lstm_inputs).toTuple();
@@ -160,14 +163,6 @@ bool PedestrianInteractionEvaluator::Evaluate(Obstacle* obstacle_ptr) {
   obstacle_id_lstm_state_map_[id].ht = ht.clone();
   obstacle_id_lstm_state_map_[id].ct = ct.clone();
 
-  std::vector<torch::jit::IValue> prediction_inputs;
-  prediction_inputs.push_back(ht[0].to(device_));
-  auto pred_out_tensor =
-      torch_prediction_layer_ptr_->forward(prediction_inputs)
-                                .toTensor()
-                                .to(torch::kCPU);
-  auto pred_out = pred_out_tensor.accessor<float, 2>();
-
   // Step 4 for-loop get a trajectory
   // Set the starting trajectory point
   Trajectory* trajectory = latest_feature_ptr->add_predicted_trajectory();
@@ -179,22 +174,62 @@ bool PedestrianInteractionEvaluator::Evaluate(Obstacle* obstacle_ptr) {
   start_point->set_v(latest_feature_ptr->speed());
   start_point->set_relative_time(0.0);
 
-  constexpr double kShortTermPredictionTimeLength = 1.0;
+  constexpr double kShortTermPredictionTimeLength = 0.5;
   int num_trajectory_point =
       static_cast<int>(kShortTermPredictionTimeLength /
                        FLAGS_prediction_trajectory_time_resolution);
-  for (int i = 1; i <= num_trajectory_point; ++i) {
+  for (int i = 1; i < num_trajectory_point; ++i) {
     double prev_x = trajectory->trajectory_point(i - 1).path_point().x();
     double prev_y = trajectory->trajectory_point(i - 1).path_point().y();
+    CHECK(obstacle_id_lstm_state_map_.find(id) !=
+          obstacle_id_lstm_state_map_.end());
+    torch::Tensor torch_position = torch::zeros({1, 2});
+    double curr_rel_x = rel_x;
+    double curr_rel_y = rel_y;
+    if (i > 1) {
+      curr_rel_x = prev_x -
+          trajectory->trajectory_point(i - 2).path_point().x();
+      curr_rel_y = prev_y -
+          trajectory->trajectory_point(i - 2).path_point().y();
+    }
+    torch_position[0][0] = curr_rel_x;
+    torch_position[0][1] = curr_rel_y;
+    std::vector<torch::jit::IValue> position_embedding_inputs;
+    position_embedding_inputs.push_back(std::move(torch_position));
+    torch::Tensor position_embedding =
+        torch_position_embedding_ptr_->forward(position_embedding_inputs)
+            .toTensor()
+            .to(torch::kCPU);
+    torch::Tensor lstm_input =
+        torch::zeros({1, kEmbeddingSize + 2 * kHiddenSize});
+    for (int i = 0; i < kEmbeddingSize; ++i) {
+      lstm_input[0][i] = position_embedding[0][i];
+    }
+    for (int i = 0; i < kHiddenSize; ++i) {
+      lstm_input[0][kEmbeddingSize + i] = ht[0][0][i];
+      lstm_input[0][kEmbeddingSize + kHiddenSize + i] = ct[0][0][i];
+    }
+    std::vector<torch::jit::IValue> lstm_inputs;
+    lstm_inputs.push_back(std::move(lstm_input));
+    auto lstm_out_tuple =
+        torch_single_lstm_ptr_->forward(lstm_inputs).toTuple();
+    ht = lstm_out_tuple->elements()[0].toTensor();
+    ct = lstm_out_tuple->elements()[1].toTensor();
+    std::vector<torch::jit::IValue> prediction_inputs;
+    prediction_inputs.push_back(ht[0]);
+    auto pred_out_tensor =
+        torch_prediction_layer_ptr_->forward(prediction_inputs)
+                                  .toTensor()
+                                  .to(torch::kCPU);
+    auto pred_out = pred_out_tensor.accessor<float, 2>();
     TrajectoryPoint* point = trajectory->add_trajectory_point();
-    double curr_x = prev_x + static_cast<double>(pred_out[0][5 * i]);
-    double curr_y = prev_y + static_cast<double>(pred_out[0][5 * i + 1]);
+    double curr_x = prev_x + static_cast<double>(pred_out[0][0]);
+    double curr_y = prev_y + static_cast<double>(pred_out[0][1]);
     point->mutable_path_point()->set_x(curr_x);
     point->mutable_path_point()->set_y(curr_y);
-    Vec2d direction(curr_x - prev_x, curr_y - prev_y);
-    double distance = direction.Length();
-    point->set_v(distance / FLAGS_prediction_trajectory_time_resolution);
-    point->mutable_path_point()->set_theta(direction.Angle());
+    point->set_v(latest_feature_ptr->speed());
+    point->mutable_path_point()->set_theta(
+        latest_feature_ptr->velocity_heading());
     point->set_relative_time(
         static_cast<double>(FLAGS_prediction_trajectory_time_resolution * i));
   }