محیط ها

با مجموعه‌ها، منظم بمانید ذخیره و دسته‌بندی محتوا براساس اولویت‌های شما.

مشاهده در TensorFlow.org در Google Colab اجرا شود مشاهده منبع در GitHub دانلود دفترچه یادداشت

معرفی

هدف یادگیری تقویتی (RL) طراحی عواملی است که از طریق تعامل با یک محیط یاد می گیرند. در تنظیمات استاندارد RL، عامل در هر مرحله زمانی یک مشاهده دریافت می کند و یک عمل را انتخاب می کند. عمل روی محیط اعمال می شود و محیط پاداش و مشاهده جدیدی را برمی گرداند. نماینده سیاستی را برای انتخاب اقداماتی برای به حداکثر رساندن مجموع پاداش ها، که به عنوان بازگشت نیز شناخته می شود، آموزش می دهد.

در TF-Agents، محیط ها را می توان در پایتون یا تنسورفلو پیاده سازی کرد. محیط‌های پایتون معمولاً پیاده‌سازی، درک و اشکال‌زدایی آسان‌تر هستند، اما محیط‌های TensorFlow کارآمدتر هستند و امکان موازی‌سازی طبیعی را فراهم می‌کنند. رایج ترین گردش کار پیاده سازی یک محیط در پایتون و استفاده از یکی از wrapper های ما برای تبدیل خودکار آن به TensorFlow است.

اجازه دهید ابتدا محیط های پایتون را بررسی کنیم. محیط های TensorFlow از یک API بسیار مشابه پیروی می کنند.

برپایی

اگر هنوز tf-agents یا gym را نصب نکرده اید، اجرا کنید:

pip install "gym>=0.21.0"
pip install tf-agents
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import tensorflow as tf
import numpy as np

from tf_agents.environments import py_environment
from tf_agents.environments import tf_environment
from tf_agents.environments import tf_py_environment
from tf_agents.environments import utils
from tf_agents.specs import array_spec
from tf_agents.environments import wrappers
from tf_agents.environments import suite_gym
from tf_agents.trajectories import time_step as ts

محیط های پایتون

محیط پایتون یک step(action) -> next_time_step روش که یک عمل به محیط زیست اعمال می شود و اطلاعات زیر در مورد گام بعدی را بر می گرداند:

  1. observation : این بخشی از دولت محیط زیست که عامل می تواند مشاهده برای انتخاب اقدامات خود را در گام بعدی است.
  2. reward : عامل است که یادگیری به حداکثر رساندن مجموع این پاداش در سراسر مراحل متعدد.
  3. step_type : تعامل با محیط زیست معمولا بخشی از یک توالی / قسمت. به عنوان مثال چندین حرکت در یک بازی شطرنج. step_type می توانند به صورت FIRST ، MID و یا LAST به نشان می دهد که آیا این گام زمان اولین، متوسط و یا آخرین مرحله در یک دنباله است.
  4. discount : این شناور به نمایندگی چقدر به وزن پاداش در زمان بعدی گام نسبت به پاداش در گام زمانی حاضر است.

این ها را به یک تاپل به نام گروه بندی می شوند TimeStep(step_type, reward, discount, observation) .

رابط کاربری که تمام محیط پایتون را باید اجرا در environments/py_environment.PyEnvironment . روش های اصلی عبارتند از:

class PyEnvironment(object):

  def reset(self):
    """Return initial_time_step."""
    self._current_time_step = self._reset()
    return self._current_time_step

  def step(self, action):
    """Apply action and return new time_step."""
    if self._current_time_step is None:
        return self.reset()
    self._current_time_step = self._step(action)
    return self._current_time_step

  def current_time_step(self):
    return self._current_time_step

  def time_step_spec(self):
    """Return time_step_spec."""

  @abc.abstractmethod
  def observation_spec(self):
    """Return observation_spec."""

  @abc.abstractmethod
  def action_spec(self):
    """Return action_spec."""

  @abc.abstractmethod
  def _reset(self):
    """Return initial_time_step."""

  @abc.abstractmethod
  def _step(self, action):
    """Apply action and return new time_step."""

علاوه بر این به step() روش، محیط نیز ارائه reset() روش که شروع می شود یک توالی جدید و اولیه فراهم می کند TimeStep . لازم نیست به پاسخ reset روش صراحت. ما فرض می‌کنیم که محیط‌ها به طور خودکار بازنشانی می‌شوند، چه زمانی که به پایان یک قسمت می‌رسند یا زمانی که step() اولین بار فراخوانی شود.

توجه داشته باشید که زیر کلاس انجام پیاده سازی نیست step() و یا reset() به طور مستقیم. در عوض، آنها را زیر پا بگذارند _step() و _reset() روش. مراحل زمان بازگشت از این روش ها می توان ذخیره سازی و در معرض از طریق current_time_step() .

observation_spec و action_spec روش بازگشت یک لانه (Bounded)ArraySpecs که به ترتیب، به نام، شکل، نوع داده و محدوده از مشاهدات و اقدامات.

در TF-Agents ما مکرراً به لانه‌هایی اشاره می‌کنیم که به‌عنوان هر ساختار درختی متشکل از لیست‌ها، تاپل‌ها، تاپل‌های نام‌گذاری شده یا دیکشنری‌ها تعریف می‌شوند. برای حفظ ساختار مشاهدات و اقدامات می‌توان این موارد را دلخواه ترکیب کرد. ما دریافتیم که این برای محیط های پیچیده تر که در آن مشاهدات و اقدامات زیادی دارید بسیار مفید است.

استفاده از محیط های استاندارد

TF عوامل ساخته شده است در لفافه برای بسیاری از محیط استاندارد مانند OpenAI بدنسازی، DeepMind کنترل و آتاری، به طوری که آنها ما را دنبال کنید py_environment.PyEnvironment رابط. این محیط های پیچیده شده را می توان به راحتی با استفاده از مجموعه های محیطی ما بارگیری کرد. بیایید محیط CartPole را از سالن ورزشی OpenAI بارگیری کنیم و به action و time_step_spec نگاه کنیم.

environment = suite_gym.load('CartPole-v0')
print('action_spec:', environment.action_spec())
print('time_step_spec.observation:', environment.time_step_spec().observation)
print('time_step_spec.step_type:', environment.time_step_spec().step_type)
print('time_step_spec.discount:', environment.time_step_spec().discount)
print('time_step_spec.reward:', environment.time_step_spec().reward)
action_spec: BoundedArraySpec(shape=(), dtype=dtype('int64'), name='action', minimum=0, maximum=1)
time_step_spec.observation: BoundedArraySpec(shape=(4,), dtype=dtype('float32'), name='observation', minimum=[-4.8000002e+00 -3.4028235e+38 -4.1887903e-01 -3.4028235e+38], maximum=[4.8000002e+00 3.4028235e+38 4.1887903e-01 3.4028235e+38])
time_step_spec.step_type: ArraySpec(shape=(), dtype=dtype('int32'), name='step_type')
time_step_spec.discount: BoundedArraySpec(shape=(), dtype=dtype('float32'), name='discount', minimum=0.0, maximum=1.0)
time_step_spec.reward: ArraySpec(shape=(), dtype=dtype('float32'), name='reward')

بنابراین می بینیم که محیط زیست انتظار اقدامات از نوع int64 در [0، 1] و بازده TimeSteps که در آن مشاهدات یک float32 بردار به طول 4 و عامل تنزیل است float32 در [0.0، 1.0. در حال حاضر، اجازه دهید سعی کنید به یک عمل ثابت (1,) برای قسمت کل.

action = np.array(1, dtype=np.int32)
time_step = environment.reset()
print(time_step)
while not time_step.is_last():
  time_step = environment.step(action)
  print(time_step)
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.0138565 , -0.03582913,  0.04861612, -0.03755046], dtype=float32),
 'reward': array(0., dtype=float32),
 'step_type': array(0, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.01313992,  0.15856317,  0.0478651 , -0.3145069 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.01631118,  0.35297176,  0.04157497, -0.5917188 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.02337062,  0.54748774,  0.02974059, -0.87102115], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.03432037,  0.74219286,  0.01232017, -1.1542072 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.04916423,  0.93715197, -0.01076398, -1.4430016 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.06790727,  1.1324048 , -0.03962401, -1.7390285 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.09055536,  1.327955  , -0.07440457, -2.04377   ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.11711447,  1.523758  , -0.11527998, -2.3585167 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.14758962,  1.7197047 , -0.16245031, -2.6843033 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(0., dtype=float32),
 'observation': array([ 0.18198372,  1.9156038 , -0.21613638, -3.0218334 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(2, dtype=int32)})

ایجاد محیط پایتون خود

برای بسیاری از مشتریان، یک مورد رایج استفاده از یکی از عوامل استاندارد (به نمایندگی ها/ مراجعه کنید) در TF-Agents برای مشکل خود استفاده می شود. برای انجام این کار، آنها باید مشکل خود را به عنوان یک محیط در نظر بگیرند. بنابراین اجازه دهید نحوه پیاده سازی یک محیط در پایتون را بررسی کنیم.

فرض کنید می‌خواهیم ماموری را برای انجام بازی کارت زیر (با الهام از جک سیاه) آموزش دهیم:

  1. بازی با استفاده از کارت های بی نهایت با شماره 1...10 انجام می شود.
  2. در هر نوبت، نماینده می تواند 2 کار انجام دهد: یک کارت تصادفی جدید دریافت کند یا دور فعلی را متوقف کند.
  3. هدف این است که مجموع کارت های خود را تا حد امکان به عدد 21 در پایان راند، بدون اینکه بیش از حد انجام دهید، به دست آورید.

محیطی که بازی را نشان می دهد می تواند به شکل زیر باشد:

  1. اقدامات: ما 2 عمل داریم. اقدام 0: یک کارت جدید دریافت کنید و اقدام 1: دور فعلی را خاتمه دهید.
  2. مشاهدات: مجموع کارت های دور فعلی.
  3. پاداش: هدف این است که تا جایی که ممکن است به عدد 21 نزدیک شویم بدون اینکه پشت سر بگذاریم، بنابراین می توانیم با استفاده از پاداش زیر در پایان راند به این کار دست یابیم: sum_of_cards - 21 اگر sum_of_cards <= 21, other -21
class CardGameEnv(py_environment.PyEnvironment):

  def __init__(self):
    self._action_spec = array_spec.BoundedArraySpec(
        shape=(), dtype=np.int32, minimum=0, maximum=1, name='action')
    self._observation_spec = array_spec.BoundedArraySpec(
        shape=(1,), dtype=np.int32, minimum=0, name='observation')
    self._state = 0
    self._episode_ended = False

  def action_spec(self):
    return self._action_spec

  def observation_spec(self):
    return self._observation_spec

  def _reset(self):
    self._state = 0
    self._episode_ended = False
    return ts.restart(np.array([self._state], dtype=np.int32))

  def _step(self, action):

    if self._episode_ended:
      # The last action ended the episode. Ignore the current action and start
      # a new episode.
      return self.reset()

    # Make sure episodes don't go on forever.
    if action == 1:
      self._episode_ended = True
    elif action == 0:
      new_card = np.random.randint(1, 11)
      self._state += new_card
    else:
      raise ValueError('`action` should be 0 or 1.')

    if self._episode_ended or self._state >= 21:
      reward = self._state - 21 if self._state <= 21 else -21
      return ts.termination(np.array([self._state], dtype=np.int32), reward)
    else:
      return ts.transition(
          np.array([self._state], dtype=np.int32), reward=0.0, discount=1.0)

بیایید مطمئن شویم که همه چیز را به درستی انجام داده ایم که محیط فوق را تعریف کرده ایم. هنگام ایجاد محیط شخصی خود، باید مطمئن شوید که مشاهدات و time_steps ایجاد شده از اشکال و انواع درست همانطور که در مشخصات شما تعریف شده است پیروی می کنند. اینها برای تولید نمودار TensorFlow استفاده می‌شوند و به همین دلیل می‌توانند مشکلاتی را برای اشکال‌زدایی در صورت اشتباه ایجاد کنند.

برای اعتبارسنجی محیط خود، از یک خط‌مشی تصادفی برای ایجاد کنش‌ها استفاده می‌کنیم و بیش از 5 قسمت را تکرار می‌کنیم تا مطمئن شویم که همه چیز همانطور که در نظر گرفته شده است کار می‌کند. اگر یک time_step دریافت کنیم که از مشخصات محیطی پیروی نمی کند، خطا ایجاد می شود.

environment = CardGameEnv()
utils.validate_py_environment(environment, episodes=5)

اکنون که می دانیم محیط طبق برنامه کار می کند، بیایید این محیط را با استفاده از یک خط مشی ثابت اجرا کنیم: درخواست 3 کارت و سپس پایان دور راند.

get_new_card_action = np.array(0, dtype=np.int32)
end_round_action = np.array(1, dtype=np.int32)

environment = CardGameEnv()
time_step = environment.reset()
print(time_step)
cumulative_reward = time_step.reward

for _ in range(3):
  time_step = environment.step(get_new_card_action)
  print(time_step)
  cumulative_reward += time_step.reward

time_step = environment.step(end_round_action)
print(time_step)
cumulative_reward += time_step.reward
print('Final Reward = ', cumulative_reward)
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([0], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(0, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([9], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([12], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(0., dtype=float32),
 'observation': array([21], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(2, dtype=int32)})
TimeStep(
{'discount': array(0., dtype=float32),
 'observation': array([21], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(2, dtype=int32)})
Final Reward =  0.0

لفاف های محیطی

یک محیط پوشاننده یک محیط پایتون را می گیرد و یک نسخه اصلاح شده از محیط را برمی گرداند. هر دو محیط اصلی و محیط زیست اصلاح موارد هستند py_environment.PyEnvironment ، و لفافه های متعدد می تواند با هم زنجیر.

برخی لفافه مشترک را می توان در یافت environments/wrappers.py . مثلا:

  1. ActionDiscretizeWrapper : تبدیل یک فضای ثابت به یک فضای گسسته.
  2. RunStats : قطاری اجرا آمار از محیط زیست از جمله تعدادی از اقدامات انجام شده، تعدادی از قسمت های تکمیل و غیره
  3. TimeLimit : پایان قسمت پس از یک عدد ثابت از مراحل.

مثال 1: Action Discretize Wrapper

InvertedPendulum یک محیط PyBullet که اقدامات مستمر در محدوده پذیرد [-2, 2] . اگر بخواهیم یک اکشن عامل گسسته مانند DQN را در این محیط آموزش دهیم، باید فضای اکشن را گسسته کنیم (کوانتیزه کنیم). این دقیقا همان چیزی است ActionDiscretizeWrapper کند. مقایسه action_spec قبل و بعد از بسته بندی:

env = suite_gym.load('Pendulum-v1')
print('Action Spec:', env.action_spec())

discrete_action_env = wrappers.ActionDiscretizeWrapper(env, num_actions=5)
print('Discretized Action Spec:', discrete_action_env.action_spec())
Action Spec: BoundedArraySpec(shape=(1,), dtype=dtype('float32'), name='action', minimum=-2.0, maximum=2.0)
Discretized Action Spec: BoundedArraySpec(shape=(), dtype=dtype('int32'), name='action', minimum=0, maximum=4)

پیچیده discrete_action_env یک نمونه از است py_environment.PyEnvironment و می تواند مانند یک محیط به طور منظم پایتون درمان می شود.

محیط های TensorFlow

رابط کاربری برای محیط های TF در تعریف environments/tf_environment.TFEnvironment و به نظر می رسد بسیار شبیه به محیط پایتون. محیط‌های TF از چند جهت با محیط‌های پایتون تفاوت دارند:

  • آنها به جای آرایه، اشیاء تانسور تولید می کنند
  • محیط های TF یک بعد دسته ای به تانسورهای تولید شده در مقایسه با مشخصات اضافه می کنند.

تبدیل محیط های پایتون به TFEnvs به tensorflow اجازه می دهد تا عملیات را موازی کند. به عنوان مثال، یک نفر می تواند تعریف collect_experience_op که جمع آوری اطلاعات از محیط زیست و می افزاید: به replay_buffer و train_op که از خواند replay_buffer و آموزش عامل، و اجرای آنها را به صورت موازی به طور طبیعی در TensorFlow.

class TFEnvironment(object):

  def time_step_spec(self):
    """Describes the `TimeStep` tensors returned by `step()`."""

  def observation_spec(self):
    """Defines the `TensorSpec` of observations provided by the environment."""

  def action_spec(self):
    """Describes the TensorSpecs of the action expected by `step(action)`."""

  def reset(self):
    """Returns the current `TimeStep` after resetting the Environment."""
    return self._reset()

  def current_time_step(self):
    """Returns the current `TimeStep`."""
    return self._current_time_step()

  def step(self, action):
    """Applies the action and returns the new `TimeStep`."""
    return self._step(action)

  @abc.abstractmethod
  def _reset(self):
    """Returns the current `TimeStep` after resetting the Environment."""

  @abc.abstractmethod
  def _current_time_step(self):
    """Returns the current `TimeStep`."""

  @abc.abstractmethod
  def _step(self, action):
    """Applies the action and returns the new `TimeStep`."""

current_time_step() روش time_step فعلی را برمی گرداند و مقدار دهی اولیه محیط زیست در صورت نیاز.

reset() نیروهای روش تنظیم مجدد در محیط زیست و بازده CURRENT_STEP.

اگر action می کند بر روی قبلی بستگی ندارد time_step tf.control_dependency در مورد نیاز است Graph حالت.

در حال حاضر، اجازه دهید ما در چگونه نگاه TFEnvironments ایجاد می کند.

ایجاد محیط TensorFlow خودتان

این پیچیده تر از ایجاد محیط در پایتون است، بنابراین در این colab آن را پوشش نمی دهیم. به عنوان مثال در دسترس است در اینجا . مورد استفاده شایع تر است به پیاده سازی محیط زیست خود را در پایتون و قرار دادن آن در TensorFlow با استفاده از ما TFPyEnvironment لفاف بسته بندی (پایین را ببینید).

پیچیدن یک محیط پایتون در TensorFlow

ما به راحتی می توانید هر محیط پایتون را به یک محیط TensorFlow با استفاده از قرار دادن TFPyEnvironment لفاف بسته بندی.

env = suite_gym.load('CartPole-v0')
tf_env = tf_py_environment.TFPyEnvironment(env)

print(isinstance(tf_env, tf_environment.TFEnvironment))
print("TimeStep Specs:", tf_env.time_step_spec())
print("Action Specs:", tf_env.action_spec())
True
TimeStep Specs: TimeStep(
{'discount': BoundedTensorSpec(shape=(), dtype=tf.float32, name='discount', minimum=array(0., dtype=float32), maximum=array(1., dtype=float32)),
 'observation': BoundedTensorSpec(shape=(4,), dtype=tf.float32, name='observation', minimum=array([-4.8000002e+00, -3.4028235e+38, -4.1887903e-01, -3.4028235e+38],
      dtype=float32), maximum=array([4.8000002e+00, 3.4028235e+38, 4.1887903e-01, 3.4028235e+38],
      dtype=float32)),
 'reward': TensorSpec(shape=(), dtype=tf.float32, name='reward'),
 'step_type': TensorSpec(shape=(), dtype=tf.int32, name='step_type')})
Action Specs: BoundedTensorSpec(shape=(), dtype=tf.int64, name='action', minimum=array(0), maximum=array(1))

توجه داشته باشید که مشخصات اکنون از نوع: (Bounded)TensorSpec .

مثال های استفاده

مثال ساده

env = suite_gym.load('CartPole-v0')

tf_env = tf_py_environment.TFPyEnvironment(env)
# reset() creates the initial time_step after resetting the environment.
time_step = tf_env.reset()
num_steps = 3
transitions = []
reward = 0
for i in range(num_steps):
  action = tf.constant([i % 2])
  # applies the action and returns the new TimeStep.
  next_time_step = tf_env.step(action)
  transitions.append([time_step, action, next_time_step])
  reward += next_time_step.reward
  time_step = next_time_step

np_transitions = tf.nest.map_structure(lambda x: x.numpy(), transitions)
print('\n'.join(map(str, np_transitions)))
print('Total reward:', reward.numpy())
[TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.0078796 , -0.04736348, -0.04966116,  0.04563603]],
      dtype=float32),
 'reward': array([0.], dtype=float32),
 'step_type': array([0], dtype=int32)}), array([0], dtype=int32), TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.00882687, -0.24173944, -0.04874843,  0.32224613]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)})]
[TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.00882687, -0.24173944, -0.04874843,  0.32224613]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)}), array([1], dtype=int32), TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.01366166, -0.04595843, -0.04230351,  0.01459712]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)})]
[TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.01366166, -0.04595843, -0.04230351,  0.01459712]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)}), array([0], dtype=int32), TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.01458083, -0.24044897, -0.04201157,  0.2936384 ]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)})]
Total reward: [3.]

کل قسمت ها

env = suite_gym.load('CartPole-v0')
tf_env = tf_py_environment.TFPyEnvironment(env)

time_step = tf_env.reset()
rewards = []
steps = []
num_episodes = 5

for _ in range(num_episodes):
  episode_reward = 0
  episode_steps = 0
  while not time_step.is_last():
    action = tf.random.uniform([1], 0, 2, dtype=tf.int32)
    time_step = tf_env.step(action)
    episode_steps += 1
    episode_reward += time_step.reward.numpy()
  rewards.append(episode_reward)
  steps.append(episode_steps)
  time_step = tf_env.reset()

num_steps = np.sum(steps)
avg_length = np.mean(steps)
avg_reward = np.mean(rewards)

print('num_episodes:', num_episodes, 'num_steps:', num_steps)
print('avg_length', avg_length, 'avg_reward:', avg_reward)
num_episodes: 5 num_steps: 131
avg_length 26.2 avg_reward: 26.2