diff --git a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java
index 987a58d3f7..4f20535ff1 100644
--- a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java
+++ b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java
@@ -477,8 +477,16 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
         if (DEBUG) {
           Log.d(TAG, "Loaded ad " + adPosition + " of " + adCount + " in group " + adGroupIndex);
         }
-        adPlaybackState = adPlaybackState.withAdCount(adGroupIndex, adCount);
-        updateAdPlaybackState();
+        int oldAdCount = adPlaybackState.adGroups[adGroupIndex].count;
+        if (adCount != oldAdCount) {
+          if (oldAdCount == C.LENGTH_UNSET) {
+            adPlaybackState = adPlaybackState.withAdCount(adGroupIndex, adCount);
+            updateAdPlaybackState();
+          } else {
+            // IMA sometimes unexpectedly decreases the ad count in an ad group.
+            Log.w(TAG, "Unexpected ad count in LOADED, " + adCount + ", expected " + oldAdCount);
+          }
+        }
         if (adGroupIndex != expectedAdGroupIndex) {
           Log.w(
               TAG,
@@ -486,6 +494,7 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
                   + expectedAdGroupIndex
                   + ", actual ad group index "
                   + adGroupIndex);
+          expectedAdGroupIndex = adGroupIndex;
         }
         break;
       case CONTENT_PAUSE_REQUESTED:
@@ -536,7 +545,7 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
     if (adsManager == null) {
       // No ads were loaded, so allow playback to start without any ads.
       pendingAdRequestContext = null;
-      adPlaybackState = new AdPlaybackState(new long[0]);
+      adPlaybackState = new AdPlaybackState();
       updateAdPlaybackState();
     }
     if (pendingAdErrorEvent == null) {
@@ -866,7 +875,7 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
         Log.d(TAG, "Unexpected CONTENT_RESUME_REQUESTED without stopAd");
       }
     }
-    if (playingAd && adGroupIndex != C.INDEX_UNSET) {
+    if (adGroupIndex != C.INDEX_UNSET) {
       adPlaybackState = adPlaybackState.withSkippedAdGroup(adGroupIndex);
       adGroupIndex = C.INDEX_UNSET;
       updateAdPlaybackState();
@@ -885,7 +894,6 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
   }
 
   private void stopAdInternal() {
-    Assertions.checkState(imaAdState != IMA_AD_STATE_NONE);
     imaAdState = IMA_AD_STATE_NONE;
     int adIndexInAdGroup = adPlaybackState.adGroups[adGroupIndex].getFirstAdIndexToPlay();
     // TODO: Handle the skipped event so the ad can be marked as skipped rather than played.
@@ -898,7 +906,9 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
   }
 
   private void handleAdGroupLoadError() {
-    AdPlaybackState.AdGroup adGroup = adPlaybackState.adGroups[expectedAdGroupIndex];
+    int adGroupIndex =
+        this.adGroupIndex == C.INDEX_UNSET ? expectedAdGroupIndex : this.adGroupIndex;
+    AdPlaybackState.AdGroup adGroup = adPlaybackState.adGroups[adGroupIndex];
     // Ad group load error can be notified more than once, so check if it was already handled.
     // TODO: Update the expected ad group index based on the position returned by
     // getContentProgress so that it's possible to detect when more than one ad group fails to load
@@ -906,10 +916,10 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
     if (adGroup.count == C.LENGTH_UNSET
         || adGroup.states[0] == AdPlaybackState.AD_STATE_UNAVAILABLE) {
       if (DEBUG) {
-        Log.d(TAG, "Removing ad group " + expectedAdGroupIndex + " as it failed to load");
+        Log.d(TAG, "Removing ad group " + adGroupIndex + " as it failed to load");
       }
-      adPlaybackState = adPlaybackState.withAdCount(expectedAdGroupIndex, 1);
-      adPlaybackState = adPlaybackState.withAdLoadError(expectedAdGroupIndex, 0);
+      adPlaybackState =
+          adPlaybackState.withAdCount(adGroupIndex, 1).withAdLoadError(adGroupIndex, 0);
       updateAdPlaybackState();
     }
   }
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java b/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java
index 5906ffb8e7..4b58a1075c 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java
@@ -495,7 +495,8 @@ public abstract class Timeline {
      * @return The duration of the ad, or {@link C#TIME_UNSET} if not yet known.
      */
     public long getAdDurationUs(int adGroupIndex, int adIndexInAdGroup) {
-      return adPlaybackState.adGroups[adGroupIndex].durationsUs[adIndexInAdGroup];
+      AdPlaybackState.AdGroup adGroup = adPlaybackState.adGroups[adGroupIndex];
+      return adGroup.count != C.LENGTH_UNSET ? adGroup.durationsUs[adIndexInAdGroup] : C.TIME_UNSET;
     }
 
     /**
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java b/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java
index 02bd67dbd8..2f6f81d6b9 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java
@@ -140,7 +140,9 @@ public final class AdPlaybackState {
       Assertions.checkArgument(count == C.LENGTH_UNSET || index < count);
       @AdState int[] states = copyStatesWithSpaceForAdCount(this.states, index + 1);
       Assertions.checkArgument(
-          states[index] == AD_STATE_UNAVAILABLE || states[index] == AD_STATE_AVAILABLE);
+          states[index] == AD_STATE_UNAVAILABLE
+              || states[index] == AD_STATE_AVAILABLE
+              || states[index] == state);
       long[] durationsUs =
           this.durationsUs.length == states.length
               ? this.durationsUs
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdsLoader.java b/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdsLoader.java
index 99feccd2f3..c2dfd91301 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdsLoader.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdsLoader.java
@@ -22,22 +22,22 @@ import java.io.IOException;
 
 /**
  * Interface for loaders of ads, which can be used with {@link AdsMediaSource}.
- * <p>
- * Ad loaders notify the {@link AdsMediaSource} about events via {@link EventListener}. In
+ *
+ * <p>Ad loaders notify the {@link AdsMediaSource} about events via {@link EventListener}. In
  * particular, implementations must call {@link EventListener#onAdPlaybackState(AdPlaybackState)}
  * with a new copy of the current {@link AdPlaybackState} whenever further information about ads
  * becomes known (for example, when an ad media URI is available, or an ad has played to the end).
- * <p>
- * {@link #attachPlayer(ExoPlayer, EventListener, ViewGroup)} will be called when the ads media
+ *
+ * <p>{@link #attachPlayer(ExoPlayer, EventListener, ViewGroup)} will be called when the ads media
  * source first initializes, at which point the loader can request ads. If the player enters the
  * background, {@link #detachPlayer()} will be called. Loaders should maintain any ad playback state
  * in preparation for a later call to {@link #attachPlayer(ExoPlayer, EventListener, ViewGroup)}. If
- * an ad is playing when the player is detached, store the current playback position via
- * {@link AdPlaybackState#setAdResumePositionUs(long)}.
- * <p>
- * If {@link EventListener#onAdPlaybackState(AdPlaybackState)} has been called, the implementation
- * of {@link #attachPlayer(ExoPlayer, EventListener, ViewGroup)} should invoke the same listener to
- * provide the existing playback state to the new player.
+ * an ad is playing when the player is detached, update the ad playback state with the current
+ * playback position using {@link AdPlaybackState#withAdResumePositionUs(long)}.
+ *
+ * <p>If {@link EventListener#onAdPlaybackState(AdPlaybackState)} has been called, the
+ * implementation of {@link #attachPlayer(ExoPlayer, EventListener, ViewGroup)} should invoke the
+ * same listener to provide the existing playback state to the new player.
  */
 public interface AdsLoader {