Trust region blackboxoptimizer algorithm partially dead code?

[boomanaiden154]

es_trainer_lib exposes this in

ml-compiler-opt/compiler_opt/es/es_trainer_lib.py

Line 164 in 2b0a74b

elif learner_config.blackbox_optimizer == (

It calls into the implementation with a bunch of flags that do not exist. This should probably be fixed (if I am understanding things correctly and this is actually an issue) or we should prune this.

[Harishankar14]

Hi @boomanaiden154 , I'd like to work on this issue. I'll investigate the trust region branch in es_trainer_lib.py and either prune the dead code or fix the missing flags, then submit a PR. If there's any additional context or if someone else is already on it,do let me know.
Happy to help ;)

[boomanaiden154]

I'm not sure there's any current action items here. I mostly filed this so that I would go back and investigate eventually. It's hard to tell what to do without additional context that is probably not all available publicly.

[Harishankar14]

@boomanaiden154 Understood, Is there any issues which i could work upon !?

[boomanaiden154]

I'm not aware of anything trivial at the moment.

[mtrofin]

@Harishankar14 depends on how much involvement/ time you have :) for example: we'd love to move the RL training to JAX.

Another issue would be moving the logging infrastructure (except tensorflow specifics) to LLVM (under utils/mlgo)

I'm sure we can find a few more.

[Harishankar14]

@mtrofin I guess, i can make some usefull contributions for moving RL training to JAX and of course more interesting ones :)

[mtrofin]

The RL -> JAX bit is quite a good chunk of independent, good work, I suspect - I'm not even sure what the state of RL algorithms is in JAX. If you're interested in it, by all means!

Created #529

[Harishankar14]

@mtrofin Thank you, will overview it and let you know :)

[Harishankar14]

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""JAX LLVM Policy Trainer."""

import time
from collections.abc import Callable
from typing import Any, Optional
import pickle
from pathlib import Path

from absl import logging
import gin
import jax
import jax.numpy as jnp
from flax import struct
from flax.training import checkpoints
from flax.training import train_state
import optax
from tensorboardX import SummaryWriter


@struct.dataclass
class TrainingMetrics:
  """Metrics for tracking training progress."""
  data_action_sum: jnp.ndarray
  data_action_count: jnp.ndarray
  data_reward_sum: jnp.ndarray
  data_reward_count: jnp.ndarray
  num_trajectories: jnp.ndarray
  correct_predictions: jnp.ndarray
  total_predictions: jnp.ndarray


@struct.dataclass
class Experience:
  """Experience data structure for trajectories."""
  observation: jnp.ndarray
  action: jnp.ndarray
  reward: jnp.ndarray
  step_type: jnp.ndarray
  discount: jnp.ndarray
  
  def is_boundary(self) -> jnp.ndarray:
    """Returns mask for boundary transitions."""
    return self.step_type == 2  # Assuming 2 indicates boundary
  
  def is_first(self) -> jnp.ndarray:
    """Returns mask for first steps."""
    return self.step_type == 0  # Assuming 0 indicates first step


def create_train_state(
    rng: jax.random.PRNGKey,
    model: Any,
    learning_rate: float,
    input_shape: tuple
) -> train_state.TrainState:
  """Creates initial training state."""
  dummy_input = jnp.ones(input_shape)
  params = model.init(rng, dummy_input)
  tx = optax.adam(learning_rate)
  return train_state.TrainState.create(
      apply_fn=model.apply,
      params=params,
      tx=tx
  )


def initialize_metrics() -> TrainingMetrics:
  """Initialize metrics with zeros."""
  return TrainingMetrics(
      data_action_sum=jnp.array(0.0),
      data_action_count=jnp.array(0.0),
      data_reward_sum=jnp.array(0.0),
      data_reward_count=jnp.array(0.0),
      num_trajectories=jnp.array(0.0),
      correct_predictions=jnp.array(0.0),
      total_predictions=jnp.array(0.0)
  )


def reset_metrics() -> TrainingMetrics:
  """Reset metrics to zero."""
  return initialize_metrics()


@jax.jit
def update_metrics(
    metrics: TrainingMetrics,
    experience: Experience,
    predicted_actions: Optional[jnp.ndarray] = None
) -> TrainingMetrics:
  """Update metrics based on experience."""
  is_action = ~experience.is_boundary()
  
  # Update action mean
  action_sum = metrics.data_action_sum + jnp.sum(
      experience.action * is_action)
  action_count = metrics.data_action_count + jnp.sum(is_action)
  
  # Update reward mean
  reward_sum = metrics.data_reward_sum + jnp.sum(
      experience.reward * is_action)
  reward_count = metrics.data_reward_count + jnp.sum(is_action)
  
  # Update trajectory count
  num_trajectories = metrics.num_trajectories + jnp.sum(experience.is_first())
  
  # Update accuracy if predictions provided
  if predicted_actions is not None:
    correct = jnp.sum((experience.action == predicted_actions) * is_action)
    total = jnp.sum(is_action)
    correct_predictions = metrics.correct_predictions + correct
    total_predictions = metrics.total_predictions + total
  else:
    correct_predictions = metrics.correct_predictions
    total_predictions = metrics.total_predictions
  
  return TrainingMetrics(
      data_action_sum=action_sum,
      data_action_count=action_count,
      data_reward_sum=reward_sum,
      data_reward_count=reward_count,
      num_trajectories=num_trajectories,
      correct_predictions=correct_predictions,
      total_predictions=total_predictions
  )


@gin.configurable
class Trainer:
  """JAX-based LLVM Policy Trainer."""

  def __init__(
      self,
      root_dir: str,
      model: Any,
      learning_rate: float = 1e-4,
      input_shape: tuple = (1, 128),
      rng_seed: int = 0,
      random_network_distillation: Any = None,
      warmstart_policy_dir: Optional[str] = None,
      checkpoint_interval: int = 10000,
      log_interval: int = 100,
      summary_log_interval: int = 100,
      summary_export_interval: int = 1000,
      summaries_flush_secs: int = 10,
      bc_percentage_correct: bool = False
  ):
    self._root_dir = Path(root_dir)
    self._root_dir.mkdir(parents=True, exist_ok=True)
    
    self._model = model
    self._random_network_distillation = random_network_distillation
    self._checkpoint_interval = checkpoint_interval
    self._log_interval = log_interval
    self._summary_log_interval = summary_log_interval
    self._summary_export_interval = summary_export_interval
    self._bc_percentage_correct = bc_percentage_correct
    
    # Initialize summary writer
    self._summary_writer = SummaryWriter(
        str(self._root_dir), 
        flush_secs=summaries_flush_secs
    )
    
    # Initialize RNG
    self._rng = jax.random.PRNGKey(rng_seed)
    
    # Initialize training state
    self._rng, init_rng = jax.random.split(self._rng)
    self._state = create_train_state(
        init_rng, model, learning_rate, input_shape
    )
    
    # Initialize global step
    self._global_step = 0
    
    # Load warmstart policy if provided
    if warmstart_policy_dir:
      self._load_warmstart_policy(warmstart_policy_dir)
    
    # Restore from checkpoint if exists
    self._restore_checkpoint()
    
    # Initialize metrics
    self._metrics = initialize_metrics()
    
    self._start_time = time.time()
    
    # JIT compile training step
    self._train_step_jit = jax.jit(self._train_step)

  def _load_warmstart_policy(self, policy_dir: str):
    policy_path = Path(policy_dir)
    if not policy_path.exists():
      raise ValueError(f"Warmstart policy directory does not exist: {policy_dir}")
    with open(policy_path / "params.pkl", "rb") as f:
      warmstart_params = pickle.load(f)
    self._state = self._state.replace(params=warmstart_params)
    logging.info(f"Loaded warmstart policy from {policy_dir}")

  def _restore_checkpoint(self):
    ckpt_dir = self._root_dir / "checkpoints"
    ckpt_dir.mkdir(exist_ok=True)
    restored = checkpoints.restore_checkpoint(
        ckpt_dir=str(ckpt_dir),
        target={"state": self._state, "global_step": self._global_step},
        prefix="checkpoint_"
    )
    if restored is not None and "state" in restored:
      self._state = restored["state"]
      self._global_step = restored["global_step"]
      logging.info(f"Restored checkpoint at step {self._global_step}")

  def _save_checkpoint(self):
    if self._global_step % self._checkpoint_interval == 0:
      ckpt_dir = self._root_dir / "checkpoints"
      checkpoints.save_checkpoint(
          ckpt_dir=str(ckpt_dir),
          target={"state": self._state, "global_step": self._global_step},
          step=self._global_step,
          prefix="checkpoint_",
          keep=3
      )
      logging.info(f"Saved checkpoint at step {self._global_step}")

  def _train_step(
      self, 
      state: train_state.TrainState, 
      experience: Experience
  ) -> tuple[train_state.TrainState, float, jnp.ndarray]:
    def loss_fn(params):
      logits = state.apply_fn(params, experience.observation)
      loss = optax.softmax_cross_entropy_with_integer_labels(
          logits, experience.action
      )
      is_action = ~experience.is_boundary()
      loss = jnp.sum(loss * is_action) / jnp.maximum(jnp.sum(is_action), 1.0)
      return loss, logits
    (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(
        state.params
    )
    state = state.apply_gradients(grads=grads)
    predicted_actions = jnp.argmax(logits, axis=-1)
    return state, loss, predicted_actions

  def _log_metrics_to_tensorboard(self, monitor_dict: dict[str, dict]):
    if self._global_step % self._summary_export_interval == 0:
      action_mean = (
          self._metrics.data_action_sum / 
          jnp.maximum(self._metrics.data_action_count, 1.0)
      )
      reward_mean = (
          self._metrics.data_reward_sum / 
          jnp.maximum(self._metrics.data_reward_count, 1.0)
      )
      self._summary_writer.add_scalar(
          'default/data_action_mean', 
          float(action_mean), 
          self._global_step
      )
      self._summary_writer.add_scalar(
          'default/data_reward_mean', 
          float(reward_mean), 
          self._global_step
      )
      self._summary_writer.add_scalar(
          'default/num_trajectories', 
          float(self._metrics.num_trajectories), 
          self._global_step
      )
      if self._bc_percentage_correct:
        accuracy = (
            self._metrics.correct_predictions / 
            jnp.maximum(self._metrics.total_predictions, 1.0)
        )
        self._summary_writer.add_scalar(
            'default/percentage_correct', 
            float(accuracy), 
            self._global_step
        )
      for name_scope, metrics_dict in monitor_dict.items():
        for key, value in metrics_dict.items():
          self._summary_writer.add_scalar(
              f'{name_scope}/{key}', 
              float(value), 
              self._global_step
          )

  def _log_experiment(self, loss: float):
    if self._global_step % self._log_interval == 0:
      logging.info(f'step = {self._global_step}, loss = {loss:.6f}')
      time_acc = time.time() - self._start_time
      steps_per_sec = self._log_interval / time_acc
      logging.info(f'{steps_per_sec:.3f} steps/sec')
      self._start_time = time.time()

  def global_step_numpy(self) -> int:
    return self._global_step

  def train(
      self,
      dataset_iter,
      monitor_dict: dict[str, dict],
      num_iterations: int,
      hooks: Optional[list[tuple[int, Callable[[], None]]]] = None
  ):
    self._metrics = reset_metrics()
    for iteration_index in range(num_iterations):
      try:
        experience = next(dataset_iter)
      except StopIteration:
        logging.warning(
            'Warning: skip training because not enough data to fill a batch.'
        )
        break
      if self._random_network_distillation:
        experience = self._random_network_distillation.train(experience)
      self._state, loss, predicted_actions = self._train_step_jit(
          self._state, experience
      )
      if self._global_step % self._summary_log_interval == 0:
        pred_actions = predicted_actions if self._bc_percentage_correct else None
        self._metrics = update_metrics(
            self._metrics, experience, pred_actions
        )
      self._global_step += 1
      self._log_metrics_to_tensorboard(monitor_dict)
      self._log_experiment(float(loss))
      self._save_checkpoint()
      if hooks is not None:
        for hook_iterations, hook_fn in hooks:
          if (iteration_index + 1) % hook_iterations == 0:
            hook_fn()
    self._summary_writer.flush()

  def save_policy(self, save_dir: str):
    save_path = Path(save_dir)
    save_path.mkdir(parents=True, exist_ok=True)
    with open(save_path / "params.pkl", "wb") as f:
      pickle.dump(self._state.params, f)
    logging.info(f"Saved policy to {save_dir}")

  def close(self):
    self._summary_writer.close()

@mtrofin Tried few things in translating it into JAX, i mean its not fully ready, many more modification have to be done, what do you think of this

file : rl/training.py

[mtrofin]

A PR is easier to review. Draft PR if work in progress.

I think the challenge with moving to JAX is dropping the dependency on tf_agents. Is there an equivalent built on JAX?

[Harishankar14]

@mtrofin Yes, indeed, it is a challenge to drop such a well documented dependency.Should look up for equivalent bult on JAX, i am assuming, there would be very less within the margin !!

What would u recommened for other approaches, other than JAX

[mtrofin]

Anything that can produce a SavedModel that can be converted to TFLite as shown in convert_saved_model.

For more context, currently, TFLite is the exchange format with the compiler when training, and the original SavedModel as the embeddable model (via tfcompile) when releasing the compiler in production. For the latter, once we're done with using EmitC, the TFLite model should become the only interface with the compiler (because we can take TFLite to TOSA and then C++ via EmitC, i.e. we won't need the SavedModel)

Even more generally, any framework that can both:

• be embedded in LLVM as an interpreter or JIT, i.e. something we can pass a model to (for training); and
• has a mechanism to convert a model to a framework-free format, like assembly or C++

would be a good candidate. If there is such an alternative, we can discuss adding the needed support in llvm, too.

[Harishankar14]

Hey @mtrofin I guess, Orbax Model would be a good choice !! What do you think

also @petrhosek thanks for the nudge on Orbax Model.
My other suggestion would be using RLib,(least priority).

[mtrofin]

We'd have to see that:

• we can express in orbax the current models we have (see the "release")
• we can train a model, at least with ES, if not also with a JAX alternative for PPO

• the scenarios described before, s/SavedModel/Orbax. In particular, for the inference part, even if there's no Orbax -> C++ or Orbax -> native code yet (which I'd expect to be the case), if there's an Orbax -> TFLite or Orbax -> TOSA conversion, it's probably reasonable to give it a try, because we are working on TOSA -> EmitC (and TFLite -> TOSA exists)

If you're interested in this, I'd recommend scoping initially to ES and doing the non-TF things separately (e.g. imagine a es-orbax next to es - fine to take deps on es for the numpy-dependent core algorithms, of course. No need to copy). In LLVM, we can have a side-by-side interpreter next to TFLite, we essentially need a specific MLModelRunner. The inference-only part, as long as we have a story that gets us to EmitC, great, I wouldn't block on that being ready.

[Harishankar14]

@mtrofin sure, will try it, and let u know through a draft pr or something, if you are okay with it !!