code: def update_obs_params(
      obs_buffer: jnp.ndarray,
      obs_mask: jnp.ndarray,
      obs_params: jnp.ndarray,
  ) -> jnp.ndarray:
      """Update observation normalization parameters."""
  
      obs_steps = obs_params[0]
      running_mean, running_var = jnp.split(obs_params[1:], 2)
      if obs_mask.ndim != obs_buffer.ndim:
          obs_mask = obs_mask.reshape(
              obs_mask.shape + (1,) * (obs_buffer.ndim - obs_mask.ndim)
          )
  
      new_steps = jnp.sum(obs_mask)
      total_steps = obs_steps + new_steps
  
      input_to_old_mean = (obs_buffer - running_mean) * obs_mask
      mean_diff = jnp.sum(input_to_old_mean / total_steps, axis=(0, 1))
      new_mean = running_mean + mean_diff
  
      input_to_new_mean = (obs_buffer - new_mean) * obs_mask
      var_diff = jnp.sum(
          input_to_new_mean * input_to_old_mean, axis=(0, 1)
      )
      new_var = running_var + var_diff
  
      return jnp.concatenate(
          [jnp.ones(1) * total_steps, new_mean, new_var]
      )
  
file_path: evojax/evojax/obs_norm.py
imports: import numpy
  from typing import Tuple
  from functools import partial
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: update_obs_params(obs_buffer, obs_mask, obs_params) and returns: jnp