001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.math.trajectory;
006
007import com.fasterxml.jackson.annotation.JsonProperty;
008import edu.wpi.first.math.geometry.Pose2d;
009import edu.wpi.first.math.geometry.Transform2d;
010import edu.wpi.first.math.trajectory.proto.TrajectoryProto;
011import edu.wpi.first.math.trajectory.proto.TrajectoryStateProto;
012import edu.wpi.first.util.protobuf.ProtobufSerializable;
013import java.util.ArrayList;
014import java.util.List;
015import java.util.Objects;
016import java.util.stream.Collectors;
017
018/**
019 * Represents a time-parameterized trajectory. The trajectory contains of various States that
020 * represent the pose, curvature, time elapsed, velocity, and acceleration at that point.
021 */
022public class Trajectory implements ProtobufSerializable {
023  /** Trajectory protobuf for serialization. */
024  public static final TrajectoryProto proto = new TrajectoryProto();
025
026  private final double m_totalTime;
027  private final List<State> m_states;
028
029  /** Constructs an empty trajectory. */
030  public Trajectory() {
031    m_states = new ArrayList<>();
032    m_totalTime = 0.0;
033  }
034
035  /**
036   * Constructs a trajectory from a vector of states.
037   *
038   * @param states A vector of states.
039   * @throws IllegalArgumentException if the vector of states is empty.
040   */
041  public Trajectory(final List<State> states) {
042    m_states = states;
043
044    if (m_states.isEmpty()) {
045      throw new IllegalArgumentException("Trajectory manually created with no states.");
046    }
047
048    m_totalTime = m_states.get(m_states.size() - 1).time;
049  }
050
051  /**
052   * Linearly interpolates between two values.
053   *
054   * @param startValue The start value.
055   * @param endValue The end value.
056   * @param t The fraction for interpolation.
057   * @return The interpolated value.
058   */
059  private static double lerp(double startValue, double endValue, double t) {
060    return startValue + (endValue - startValue) * t;
061  }
062
063  /**
064   * Linearly interpolates between two poses.
065   *
066   * @param startValue The start pose.
067   * @param endValue The end pose.
068   * @param t The fraction for interpolation.
069   * @return The interpolated pose.
070   */
071  private static Pose2d lerp(Pose2d startValue, Pose2d endValue, double t) {
072    return startValue.plus((endValue.minus(startValue)).times(t));
073  }
074
075  /**
076   * Returns the initial pose of the trajectory.
077   *
078   * @return The initial pose of the trajectory.
079   */
080  public Pose2d getInitialPose() {
081    return sample(0).pose;
082  }
083
084  /**
085   * Returns the overall duration of the trajectory.
086   *
087   * @return The duration of the trajectory in seconds.
088   */
089  public double getTotalTime() {
090    return m_totalTime;
091  }
092
093  /**
094   * Return the states of the trajectory.
095   *
096   * @return The states of the trajectory.
097   */
098  public List<State> getStates() {
099    return m_states;
100  }
101
102  /**
103   * Sample the trajectory at a point in time.
104   *
105   * @param time The point in time since the beginning of the trajectory to sample in seconds.
106   * @return The state at that point in time.
107   * @throws IllegalStateException if the trajectory has no states.
108   */
109  public State sample(double time) {
110    if (m_states.isEmpty()) {
111      throw new IllegalStateException("Trajectory cannot be sampled if it has no states.");
112    }
113
114    if (time <= m_states.get(0).time) {
115      return m_states.get(0);
116    }
117    if (time >= m_totalTime) {
118      return m_states.get(m_states.size() - 1);
119    }
120
121    // To get the element that we want, we will use a binary search algorithm
122    // instead of iterating over a for-loop. A binary search is O(std::log(n))
123    // whereas searching using a loop is O(n).
124
125    // This starts at 1 because we use the previous state later on for
126    // interpolation.
127    int low = 1;
128    int high = m_states.size() - 1;
129
130    while (low != high) {
131      int mid = (low + high) / 2;
132      if (m_states.get(mid).time < time) {
133        // This index and everything under it are less than the requested
134        // timestamp. Therefore, we can discard them.
135        low = mid + 1;
136      } else {
137        // t is at least as large as the element at this index. This means that
138        // anything after it cannot be what we are looking for.
139        high = mid;
140      }
141    }
142
143    // High and Low should be the same.
144
145    // The sample's timestamp is now greater than or equal to the requested
146    // timestamp. If it is greater, we need to interpolate between the
147    // previous state and the current state to get the exact state that we
148    // want.
149    final State sample = m_states.get(low);
150    final State prevSample = m_states.get(low - 1);
151
152    // If the difference in states is negligible, then we are spot on!
153    if (Math.abs(sample.time - prevSample.time) < 1E-9) {
154      return sample;
155    }
156    // Interpolate between the two states for the state that we want.
157    return prevSample.interpolate(
158        sample, (time - prevSample.time) / (sample.time - prevSample.time));
159  }
160
161  /**
162   * Transforms all poses in the trajectory by the given transform. This is useful for converting a
163   * robot-relative trajectory into a field-relative trajectory. This works with respect to the
164   * first pose in the trajectory.
165   *
166   * @param transform The transform to transform the trajectory by.
167   * @return The transformed trajectory.
168   */
169  public Trajectory transformBy(Transform2d transform) {
170    var firstState = m_states.get(0);
171    var firstPose = firstState.pose;
172
173    // Calculate the transformed first pose.
174    var newFirstPose = firstPose.plus(transform);
175    List<State> newStates = new ArrayList<>();
176
177    newStates.add(
178        new State(
179            firstState.time,
180            firstState.velocity,
181            firstState.acceleration,
182            newFirstPose,
183            firstState.curvature));
184
185    for (int i = 1; i < m_states.size(); i++) {
186      var state = m_states.get(i);
187      // We are transforming relative to the coordinate frame of the new initial pose.
188      newStates.add(
189          new State(
190              state.time,
191              state.velocity,
192              state.acceleration,
193              newFirstPose.plus(state.pose.minus(firstPose)),
194              state.curvature));
195    }
196
197    return new Trajectory(newStates);
198  }
199
200  /**
201   * Transforms all poses in the trajectory so that they are relative to the given pose. This is
202   * useful for converting a field-relative trajectory into a robot-relative trajectory.
203   *
204   * @param pose The pose that is the origin of the coordinate frame that the current trajectory
205   *     will be transformed into.
206   * @return The transformed trajectory.
207   */
208  public Trajectory relativeTo(Pose2d pose) {
209    return new Trajectory(
210        m_states.stream()
211            .map(
212                state ->
213                    new State(
214                        state.time,
215                        state.velocity,
216                        state.acceleration,
217                        state.pose.relativeTo(pose),
218                        state.curvature))
219            .collect(Collectors.toList()));
220  }
221
222  /**
223   * Concatenates another trajectory to the current trajectory. The user is responsible for making
224   * sure that the end pose of this trajectory and the start pose of the other trajectory match (if
225   * that is the desired behavior).
226   *
227   * @param other The trajectory to concatenate.
228   * @return The concatenated trajectory.
229   */
230  public Trajectory concatenate(Trajectory other) {
231    // If this is a default constructed trajectory with no states, then we can
232    // simply return the rhs trajectory.
233    if (m_states.isEmpty()) {
234      return other;
235    }
236
237    // Deep copy the current states.
238    List<State> states =
239        m_states.stream()
240            .map(
241                state ->
242                    new State(
243                        state.time,
244                        state.velocity,
245                        state.acceleration,
246                        state.pose,
247                        state.curvature))
248            .collect(Collectors.toList());
249
250    // Here we omit the first state of the other trajectory because we don't want
251    // two time points with different states. Sample() will automatically
252    // interpolate between the end of this trajectory and the second state of the
253    // other trajectory.
254    for (int i = 1; i < other.getStates().size(); ++i) {
255      var s = other.getStates().get(i);
256      states.add(new State(s.time + m_totalTime, s.velocity, s.acceleration, s.pose, s.curvature));
257    }
258    return new Trajectory(states);
259  }
260
261  /**
262   * Represents a time-parameterized trajectory. The trajectory contains of various States that
263   * represent the pose, curvature, time elapsed, velocity, and acceleration at that point.
264   */
265  public static class State implements ProtobufSerializable {
266    /** Trajectory.State protobuf for serialization. */
267    public static final TrajectoryStateProto proto = new TrajectoryStateProto();
268
269    /** The time elapsed since the beginning of the trajectory in seconds. */
270    @JsonProperty("time")
271    public double time;
272
273    /** The speed at that point of the trajectory in meters per second. */
274    @JsonProperty("velocity")
275    public double velocity;
276
277    /** The acceleration at that point of the trajectory in m/s². */
278    @JsonProperty("acceleration")
279    public double acceleration;
280
281    /** The pose at that point of the trajectory. */
282    @JsonProperty("pose")
283    public Pose2d pose;
284
285    /** The curvature at that point of the trajectory in rad/m. */
286    @JsonProperty("curvature")
287    public double curvature;
288
289    /** Default constructor. */
290    public State() {
291      pose = Pose2d.kZero;
292    }
293
294    /**
295     * Constructs a State with the specified parameters.
296     *
297     * @param time The time elapsed since the beginning of the trajectory in seconds.
298     * @param velocity The speed at that point of the trajectory in m/s.
299     * @param acceleration The acceleration at that point of the trajectory in m/s².
300     * @param pose The pose at that point of the trajectory.
301     * @param curvature The curvature at that point of the trajectory in rad/m.
302     */
303    public State(double time, double velocity, double acceleration, Pose2d pose, double curvature) {
304      this.time = time;
305      this.velocity = velocity;
306      this.acceleration = acceleration;
307      this.pose = pose;
308      this.curvature = curvature;
309    }
310
311    /**
312     * Interpolates between two States.
313     *
314     * @param endValue The end value for the interpolation.
315     * @param i The interpolant (fraction).
316     * @return The interpolated state.
317     */
318    State interpolate(State endValue, double i) {
319      // Find the new t value.
320      final double newT = lerp(time, endValue.time, i);
321
322      // Find the delta time between the current state and the interpolated state.
323      final double deltaT = newT - time;
324
325      // If delta time is negative, flip the order of interpolation.
326      if (deltaT < 0) {
327        return endValue.interpolate(this, 1 - i);
328      }
329
330      // Check whether the robot is reversing at this stage.
331      final boolean reversing = velocity < 0 || Math.abs(velocity) < 1E-9 && acceleration < 0;
332
333      // Calculate the new velocity
334      // v_f = v_0 + at
335      final double newV = velocity + (acceleration * deltaT);
336
337      // Calculate the change in position.
338      // delta_s = v_0 t + 0.5at²
339      final double newS =
340          (velocity * deltaT + 0.5 * acceleration * Math.pow(deltaT, 2)) * (reversing ? -1.0 : 1.0);
341
342      // Return the new state. To find the new position for the new state, we need
343      // to interpolate between the two endpoint poses. The fraction for
344      // interpolation is the change in position (delta s) divided by the total
345      // distance between the two endpoints.
346      final double interpolationFrac =
347          newS / endValue.pose.getTranslation().getDistance(pose.getTranslation());
348
349      return new State(
350          newT,
351          newV,
352          acceleration,
353          lerp(pose, endValue.pose, interpolationFrac),
354          lerp(curvature, endValue.curvature, interpolationFrac));
355    }
356
357    @Override
358    public String toString() {
359      return String.format(
360          "State(Sec: %.2f, Vel m/s: %.2f, Accel m/s/s: %.2f, Pose: %s, Curvature: %.2f)",
361          time, velocity, acceleration, pose, curvature);
362    }
363
364    @Override
365    public boolean equals(Object obj) {
366      return obj instanceof State state
367          && Double.compare(state.time, time) == 0
368          && Double.compare(state.velocity, velocity) == 0
369          && Double.compare(state.acceleration, acceleration) == 0
370          && Double.compare(state.curvature, curvature) == 0
371          && Objects.equals(pose, state.pose);
372    }
373
374    @Override
375    public int hashCode() {
376      return Objects.hash(time, velocity, acceleration, pose, curvature);
377    }
378  }
379
380  @Override
381  public String toString() {
382    String stateList = m_states.stream().map(State::toString).collect(Collectors.joining(", \n"));
383    return String.format("Trajectory - Seconds: %.2f, States:\n%s", m_totalTime, stateList);
384  }
385
386  @Override
387  public int hashCode() {
388    return m_states.hashCode();
389  }
390
391  @Override
392  public boolean equals(Object obj) {
393    return obj instanceof Trajectory other && m_states.equals(other.getStates());
394  }
395}