Remove shortening logic for GAIL when demo small

Update typing

ml-agents/mlagents/trainers/components/reward_signals/curiosity/signal.py

@@ -42,7 +42,7 @@ def __init__(
  }
  self.has_updated = False
 
- def evaluate_batch(self, mini_batch: Dict[str, np.array]) -> RewardSignalResult:
+ def evaluate_batch(self, mini_batch: AgentBuffer) -> RewardSignalResult:
  feed_dict: Dict[tf.Tensor, Any] = {
  self.policy.batch_size_ph: len(mini_batch["actions"]),
  self.policy.sequence_length_ph: self.policy.sequence_length,

ml-agents/mlagents/trainers/components/reward_signals/extrinsic/signal.py

@@ -2,6 +2,7 @@
 import numpy as np
 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
+from mlagents.trainers.buffer import AgentBuffer
 
 
 class ExtrinsicRewardSignal(RewardSignal):
@@ -16,6 +17,6 @@ def check_config(
  param_keys = ["strength", "gamma"]
  super().check_config(config_dict, param_keys)
 
- def evaluate_batch(self, mini_batch: Dict[str, np.array]) -> RewardSignalResult:
+ def evaluate_batch(self, mini_batch: AgentBuffer) -> RewardSignalResult:
  env_rews = np.array(mini_batch["environment_rewards"], dtype=np.float32)
  return RewardSignalResult(self.strength * env_rews, env_rews)

ml-agents/mlagents/trainers/components/reward_signals/gail/signal.py

@@ -62,7 +62,7 @@ def __init__(
  "Policy/GAIL Expert Estimate": "gail_expert_estimate",
  }
 
- def evaluate_batch(self, mini_batch: Dict[str, np.array]) -> RewardSignalResult:
+ def evaluate_batch(self, mini_batch: AgentBuffer) -> RewardSignalResult:
  feed_dict: Dict[tf.Tensor, Any] = {
  self.policy.batch_size_ph: len(mini_batch["actions"]),
  self.policy.sequence_length_ph: self.policy.sequence_length,
@@ -110,16 +110,9 @@ def prepare_update(
  :param mini_batch_policy: A mini batch of trajectories sampled from the current policy
  :return: Feed_dict for update process.
  """
- max_num_experiences = min(
- len(mini_batch["actions"]), self.demonstration_buffer.num_experiences
- )
- # If num_sequences is less, we need to shorten the input batch.
- for key, element in mini_batch.items():
- mini_batch[key] = element[:max_num_experiences]
-
- # Get batch from demo buffer
+ # Get batch from demo buffer. Even if demo buffer is smaller, we sample with replacement
  mini_batch_demo = self.demonstration_buffer.sample_mini_batch(
- len(mini_batch["actions"]), 1
+ mini_batch.num_experiences, 1
  )
 
  feed_dict: Dict[tf.Tensor, Any] = {