class_docstr: This policy ignores inputs and is specific to a certain MDKP config.
class_docstr_tok: ["This","policy","ignores","inputs","and","is","specific","to","a","certain","MDKP","config","."]
code: def __init__(self, num_items, num_bins):
      self.num_params = num_items
      thresholds = jnp.ones([num_bins + 1, num_items]) / (num_bins + 1)
      thresholds = jnp.cumsum(thresholds, axis=0)
      thresholds = thresholds - thresholds.min(axis=0)
  
      def forward_fn(params, obs):
          if num_bins == 1:
              diff = jax.nn.sigmoid(params)[None, :] - 0.5
          else:
              diff = (
                  jnp.clip(params - 0.04, a_min=0, a_max=1)[None, :]
                  - thresholds
              )
          selection = jnp.where(diff >= 0, 1, 0).sum(axis=0) - 1
          act = (
              jnp.zeros([num_bins + 1, num_items])
              .at[selection, jnp.arange(num_items)]
              .set(1)
          )
          return act[1:]
  
      self._forward_fn = jax.vmap(forward_fn)
  
file_path: evojax/examples/train_mdkp.py
imports: import argparse
  import os
  import sys
  import numpy
  from typing import Tuple
  import jax
  import jax.numpy
  from evojax import Trainer
  from evojax.task.mdkp import MDKP
  from evojax.task.mdkp import TaskState
  from evojax.policy import PolicyNetwork
  from evojax.policy.base import PolicyState
  from evojax.algo import SimpleGA
  from evojax import util
label_desc: Use image node_id 1 to create a new MDKPPolicy object from inherited base classes: PolicyNetwork with example: obj = MDKPPolicy(num_items, num_bins)